Add evaluation code

.gitignore

@@ -8,3 +8,6 @@ wheels/
 
 # Virtual environments
 .venv
+
+# Benchmarks
+benchmarks/

README.md

@@ -19,15 +19,10 @@ With [uv](https://docs.astral.sh/uv/), add to your project with:
 uv add cells2table
 ```
 
-ONNX models need a [ONNX Runtime](https://onnxruntime.ai/getting-started) installed to run. You can install one on your own or use one of the optionals already configured.
-
 | Optional        | Description             |
 | --------------- | ----------------------- |
 | `docling`       | For docling usage       |
 | `huggingface`   | For downloading models  |
-| `onnx_cuda`     | For NVIDIA GPUs         |
-| `onnx_openvino` | For Intel GPUs and CPUs |
-| `onnx_cpu`      | Default CPU runtime     |
 
 ## Usage
 
@@ -50,7 +45,7 @@ pipeline_options = PdfPipelineOptions(
 converter = DocumentConverter(
     format_options={
         InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
-        InputFormat.IMAGE: PdfFormatOption(pipeline_options=pipeline_options),
+        InputFormat.IMAGE: ImageFormatOption(pipeline_options=pipeline_options),
     }
 )
 

cells2table/cli.py

@@ -5,12 +5,12 @@
 import cv2
 
 from cells2table.pipelines import DefaultPipeline
-from cells2table.utils.visualize import visualize_table
+from cells2table.utils.visualize import bgr_to_rgb, rgb_to_bgr, show_image, visualize_table
 
 logger = logging.getLogger(__name__)
 
 
-def download(local_dir: Path | str | None = None) -> None:
+def download() -> None:
     """Download default pipeline models."""
 
     log_format = "%(asctime)s\t%(levelname)s\t%(name)s: %(message)s"
@@ -39,11 +39,12 @@ def main() -> None:
     if not args.image_path.exists():
         raise FileNotFoundError(f"File does not exist: {args.image_path}")
 
-    image = cv2.imread(str(args.image_path))
-
+    image = cv2.imread(args.image_path)
     if image is None:
         raise ValueError(f"Failed to load image: {args.image_path}")
 
+    image = bgr_to_rgb(image)  # ty:ignore[invalid-argument-type]
+
     logger.info("Image loaded successfully from %s", args.image_path)
     logger.debug(
         "Image proprieties: width=%d, height=%d, channels=%d, datatype=%s",
@@ -57,7 +58,7 @@ def main() -> None:
     tables = table_pipeline([image])
 
     for table in tables:
-        visualize_table(image, table)  # type: ignore
+        show_image(visualize_table(rgb_to_bgr(image), table))
 
 
 if __name__ == "__main__":

cells2table/datamodels/table.py

@@ -16,6 +16,22 @@ class Cell:
     col_span: int = 1
 
 
+def sort_cells_index_by_top(cells: list[Cell]) -> list[int]:
+    return sorted(range(len(cells)), key=lambda i: cells[i].bbox.t)
+
+
+def sort_cells_index_by_bottom(cells: list[Cell]) -> list[int]:
+    return sorted(range(len(cells)), key=lambda i: cells[i].bbox.b)
+
+
+def sort_cells_index_by_left(cells: list[Cell]) -> list[int]:
+    return sorted(range(len(cells)), key=lambda i: cells[i].bbox.l)
+
+
+def sort_cells_index_by_right(cells: list[Cell]) -> list[int]:
+    return sorted(range(len(cells)), key=lambda i: cells[i].bbox.r)
+
+
 @dataclass
 class Table:
     cells: list[Cell] = field(default_factory=list)
@@ -31,87 +47,87 @@ def from_detections(cls, cells_det: Iterable[DetectionResult], tolerance: float
             cell = Cell(bbox=bbox, row=0, col=0)
             table.cells.append(cell)
 
-        table.compute_rows_and_cols(tolerance)
+        table.compute_structure(tolerance)
         return table
 
-    def compute_rows_and_cols(self, tolerance: float) -> None:
-        self.compute_rows(tolerance)
-        self.compute_cols(tolerance)
-
-    def sort_cells_by_rows(self) -> None:
-        self.cells = sorted(self.cells, key=lambda cell: cell.bbox.t)
-
-    def sort_cells_by_cols(self) -> None:
-        self.cells = sorted(self.cells, key=lambda cell: cell.bbox.l)
-
-    def compute_rows(self, tolerance: float) -> None:
-        self.sort_cells_by_rows()
-
-        row_y = None
-        row_num = 0
-        row_start_idx = 0
-        row_end_idx = None
-        check_span_indices: set[int] = set()
-
-        for i in range(len(self.cells)):
-            if row_y is None:
-                row_y = self.cells[i].bbox.t
-
-            elif abs(self.cells[i].bbox.t - row_y) > tolerance:
-                row_end_idx = i
-                for j in range(row_start_idx, row_end_idx):
-                    self.cells[j].row = row_num
-                    check_span_indices.add(j)
-
-                row_y = self.cells[i].bbox.t
-                row_start_idx = row_end_idx
-                row_num += 1
-
-                for j in list(check_span_indices):
-                    if self.cells[j].bbox.b > row_y + tolerance:
-                        self.cells[j].row_span += 1
-                    else:
-                        check_span_indices.remove(j)
-
-        row_end_idx = len(self.cells)
-        for j in range(row_start_idx, row_end_idx):
-            self.cells[j].row = row_num
-            check_span_indices.add(j)
-
-        self.num_rows = row_num + 1
-
-    def compute_cols(self, tolerance: float) -> None:
-        self.sort_cells_by_cols()
-
-        col_x = None
-        col_num = 0
-        col_start_idx = 0
-        col_end_idx = None
-        check_span_indices: set[int] = set()
-
-        for i in range(len(self.cells)):
-            if col_x is None:
-                col_x = self.cells[i].bbox.l
-
-            elif abs(self.cells[i].bbox.l - col_x) > tolerance:
-                col_end_idx = i
-                for j in range(col_start_idx, col_end_idx):
-                    self.cells[j].col = col_num
-                    check_span_indices.add(j)
-
-                col_x = self.cells[i].bbox.l
-                col_start_idx = col_end_idx
-                col_num += 1
-
-                for j in list(check_span_indices):
-                    if self.cells[j].bbox.r > col_x + tolerance:
-                        self.cells[j].col_span += 1
-                    else:
-                        check_span_indices.remove(j)
-
-        col_end_idx = len(self.cells)
-        for j in range(col_start_idx, col_end_idx):
-            self.cells[j].col = col_num
-            check_span_indices.add(j)
-
-        self.num_cols = col_num + 1
+    def compute_structure(self, tolerance: float) -> None:
+        self.compute_cells_row(tolerance)
+        self.compute_cells_col(tolerance)
+        self.compute_cells_row_span(tolerance)
+        self.compute_cells_col_span(tolerance)
+
+    def compute_cells_row(self, tolerance: float) -> None:
+        indices = sort_cells_index_by_top(self.cells)
+
+        spos = None  # Spatial position
+        lpos = 0  # Logical position
+
+        for i in indices:
+            cell_spos = self.cells[i].bbox.t
+
+            if spos is None:
+                spos = cell_spos
+
+            elif cell_spos - spos > tolerance:
+                lpos += 1
+                spos = cell_spos
+
+            self.cells[i].row = lpos
+
+        self.num_rows = lpos + 1
+
+    def compute_cells_col(self, tolerance: float) -> None:
+        indices = sort_cells_index_by_left(self.cells)
+
+        spos = None  # Spatial position
+        lpos = 0  # Logical position
+
+        for i in indices:
+            cell_spos = self.cells[i].bbox.l
+
+            if spos is None:
+                spos = cell_spos
+
+            elif cell_spos - spos > tolerance:
+                lpos += 1
+                spos = cell_spos
+
+            self.cells[i].col = lpos
+
+        self.num_cols = lpos + 1
+
+    def compute_cells_row_span(self, tolerance: float) -> None:
+        indices = sort_cells_index_by_bottom(self.cells)
+
+        spos = None  # Spatial position
+        lpos = 0  # Logical position
+
+        for i in indices:
+            cell_spos = self.cells[i].bbox.b
+
+            if spos is None:
+                spos = cell_spos
+
+            elif cell_spos - spos > tolerance:
+                lpos += 1
+                spos = cell_spos
+
+            self.cells[i].row_span = 1 + lpos - self.cells[i].row
+
+    def compute_cells_col_span(self, tolerance: float) -> None:
+        indices = sort_cells_index_by_right(self.cells)
+
+        spos = None  # Spatial position
+        lpos = 0  # Logical position
+
+        for i in indices:
+            cell_spos = self.cells[i].bbox.r
+
+            if spos is None:
+                spos = cell_spos
+
+            elif cell_spos - spos > tolerance:
+                lpos += 1
+                spos = cell_spos
+
+            self.cells[i].col_span = 1 + lpos - self.cells[i].col