cells2table 0.2.0__tar.gz

cells2table-0.2.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 João Speranza Pastorello
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

cells2table-0.2.0/PKG-INFO

@@ -0,0 +1,85 @@
+Metadata-Version: 2.4
+Name: cells2table
+Version: 0.2.0
+Summary: Table image parsing with cell detection models
+Keywords: docling,plugin
+Author: jspast
+Author-email: jspast <140563347+jspast@users.noreply.github.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Dist: huggingface-hub>=0.36.0
+Requires-Dist: numpy>=2.2.0
+Requires-Dist: opencv-python>=4.11.0.86
+Requires-Dist: onnxruntime>=1.23.2 ; extra == 'cpu'
+Requires-Dist: onnxruntime-gpu>=1.23.2 ; extra == 'cuda'
+Requires-Dist: docling>=2.66.0 ; extra == 'docling'
+Requires-Dist: onnxruntime-openvino>=1.23.0 ; extra == 'openvino'
+Requires-Python: >=3.12
+Project-URL: Homepage, https://github.com/jspast/cells2table
+Project-URL: Issues, https://github.com/jspast/cells2table/issues
+Provides-Extra: cpu
+Provides-Extra: cuda
+Provides-Extra: docling
+Provides-Extra: openvino
+Description-Content-Type: text/markdown
+
+# cells2table
+
+Parsing tables in document images with cell detection models
+
+## Implemented pipelines
+
+### PaddlePaddle models
+
+- Classification model (wired / wireless)
+- Cell detection model with different weights for each class
+
+Using [ONNX weights](https://huggingface.co/jspast/paddlepaddle-table-models-onnx) (downloaded automatically on first use with `huggingface_hub`)
+
+## Instalation
+
+With [uv](https://docs.astral.sh/uv/), add to your project with:
+
+```sh
+uv add git+https://github.com/jspast/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             |
+| ---------- | ----------------------- |
+| `cuda`     | For NVIDIA GPUs         |
+| `openvino` | For Intel GPUs and CPUs |
+| `cpu`      | Default CPU runtime     |
+| `docling`  | For docling usage       |
+
+## Usage
+
+cells2table only extract structural information from the tables. Another library is needed to extract content from the cells.
+
+### Docling
+
+A [docling plugin](https://docling-project.github.io/docling/concepts/plugins/) is provided to allow integrating cells2table in a complete pipeline.
+
+Usage example:
+
+```python
+from cells2table.docling import CustomDoclingTableStructureOptions
+
+pipeline_options = PdfPipelineOptions(
+    allow_external_plugins=True,
+    table_structure_options=CustomDoclingTableStructureOptions(),
+)
+
+converter = DocumentConverter(
+    format_options={
+        InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
+        InputFormat.IMAGE: PdfFormatOption(pipeline_options=pipeline_options),
+    }
+)
+
+result = converter.convert("path/to/document.pdf")
+print(result.document.export_to_markdown())
+```

cells2table-0.2.0/README.md

@@ -0,0 +1,58 @@
+# cells2table
+
+Parsing tables in document images with cell detection models
+
+## Implemented pipelines
+
+### PaddlePaddle models
+
+- Classification model (wired / wireless)
+- Cell detection model with different weights for each class
+
+Using [ONNX weights](https://huggingface.co/jspast/paddlepaddle-table-models-onnx) (downloaded automatically on first use with `huggingface_hub`)
+
+## Instalation
+
+With [uv](https://docs.astral.sh/uv/), add to your project with:
+
+```sh
+uv add git+https://github.com/jspast/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             |
+| ---------- | ----------------------- |
+| `cuda`     | For NVIDIA GPUs         |
+| `openvino` | For Intel GPUs and CPUs |
+| `cpu`      | Default CPU runtime     |
+| `docling`  | For docling usage       |
+
+## Usage
+
+cells2table only extract structural information from the tables. Another library is needed to extract content from the cells.
+
+### Docling
+
+A [docling plugin](https://docling-project.github.io/docling/concepts/plugins/) is provided to allow integrating cells2table in a complete pipeline.
+
+Usage example:
+
+```python
+from cells2table.docling import CustomDoclingTableStructureOptions
+
+pipeline_options = PdfPipelineOptions(
+    allow_external_plugins=True,
+    table_structure_options=CustomDoclingTableStructureOptions(),
+)
+
+converter = DocumentConverter(
+    format_options={
+        InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
+        InputFormat.IMAGE: PdfFormatOption(pipeline_options=pipeline_options),
+    }
+)
+
+result = converter.convert("path/to/document.pdf")
+print(result.document.export_to_markdown())
+```

cells2table-0.2.0/pyproject.toml

@@ -0,0 +1,56 @@
+[project]
+name = "cells2table"
+version = "0.2.0"
+description = "Table image parsing with cell detection models"
+readme = "README.md"
+requires-python = ">=3.12"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [
+  { name = "jspast", email = "140563347+jspast@users.noreply.github.com" }
+]
+keywords = ["docling", "plugin"]
+classifiers = [
+  "Operating System :: OS Independent",
+  "Programming Language :: Python :: 3",
+]
+dependencies = [
+  "huggingface-hub>=0.36.0",
+  "numpy>=2.2.0",
+  "opencv-python>=4.11.0.86",
+]
+
+[project.urls]
+Homepage = "https://github.com/jspast/cells2table"
+Issues = "https://github.com/jspast/cells2table/issues"
+
+[project.scripts]
+cells2table = "cells2table.cli:main"
+
+[project.entry-points.docling]
+cells2table = "cells2table.docling"
+
+[project.optional-dependencies]
+cpu = [
+  "onnxruntime>=1.23.2",
+]
+cuda = [
+  "onnxruntime-gpu>=1.23.2",
+]
+docling = [
+  "docling>=2.66.0",
+]
+openvino = [
+  "onnxruntime-openvino>=1.23.0",
+]
+
+[build-system]
+requires = ["uv_build>=0.9.17,<0.10.0"]
+build-backend = "uv_build"
+
+[tool.ruff]
+line-length = 100
+
+[tool.uv]
+package = true
+conflicts = [[{ extra = "cpu" }, { extra = "cuda" }, { extra = "openvino" }]]

cells2table-0.2.0/src/cells2table/__init__.py

@@ -0,0 +1,3 @@
+from .models import DefaultPipeline
+
+__all__ = [DefaultPipeline]

cells2table-0.2.0/src/cells2table/docling.py

@@ -0,0 +1,270 @@
+import copy
+import logging
+from collections.abc import Iterable
+from pathlib import Path
+from typing import ClassVar, Literal, Optional, Sequence, Type
+
+import numpy
+from docling.datamodel.accelerator_options import AcceleratorOptions
+from docling.datamodel.base_models import Cluster, Page, Table, TableStructurePrediction
+from docling.datamodel.document import ConversionResult
+from docling.datamodel.pipeline_options import BaseTableStructureOptions
+from docling.datamodel.settings import settings
+from docling.models.base_table_model import BaseTableStructureModel
+from docling.utils.profiling import TimeRecorder
+from docling_core.types.doc.base import BoundingBox
+from docling_core.types.doc.document import TableCell
+from docling_core.types.doc.labels import DocItemLabel
+from docling_core.types.doc.page import BoundingRectangle, TextCellUnit
+from PIL import ImageDraw
+
+# from docling.utils.accelerator_utils import decide_device
+from .models import DefaultPipeline
+
+logger = logging.getLogger(__name__)
+
+
+def get_tokens(page: Page, table_cluster: Cluster, scale: float) -> list[str]:
+    """Docling logic for token extraction."""
+
+    # Check if word-level cells are available from backend:
+    sp = page._backend.get_segmented_page() if page._backend else None
+    if sp is not None:
+        tcells = sp.get_cells_in_bbox(
+            cell_unit=TextCellUnit.WORD,
+            bbox=table_cluster.bbox,
+        )
+        if len(tcells) == 0:
+            # In case word-level cells yield empty
+            tcells = table_cluster.cells
+    else:
+        # Otherwise - we use normal (line/phrase) cells
+        tcells = table_cluster.cells
+    tokens = []
+    for c in tcells:
+        # Only allow non empty strings (spaces) into the cells of a table
+        if len(c.text.strip()) > 0:
+            new_cell = copy.deepcopy(c)
+            new_cell.rect = BoundingRectangle.from_bounding_box(
+                new_cell.rect.to_bounding_box().scaled(scale=scale)
+            )
+            tokens.append(
+                {
+                    "id": new_cell.index,
+                    "text": new_cell.text,
+                    "bbox": new_cell.rect.to_bounding_box().model_dump(),
+                }
+            )
+
+    return tokens
+
+
+class CustomDoclingTableStructureOptions(BaseTableStructureOptions):
+    kind: ClassVar[Literal["cells2table"]] = "cells2table"
+
+
+class CustomDoclingTableStructureModel(BaseTableStructureModel):
+    def __init__(
+        self,
+        enabled: bool,
+        artifacts_path: Optional[Path],
+        options: CustomDoclingTableStructureOptions,
+        accelerator_options: AcceleratorOptions,
+    ):
+        self.enabled = enabled
+
+        if self.enabled:
+            self.pipeline = DefaultPipeline(artifacts_path)
+
+            # TODO: decide how to deal with accelerator options
+            # device = decide_device(accelerator_options.device)
+
+            self.scale = 2.0  # Scale up table input images to 144 dpi
+
+    @classmethod
+    def get_options_type(cls) -> Type[BaseTableStructureOptions]:
+        return CustomDoclingTableStructureOptions
+
+    def predict_tables(
+        self,
+        conv_res: ConversionResult,
+        pages: Sequence[Page],
+    ) -> Sequence[TableStructurePrediction]:
+        pages = list(pages)
+        predictions: list[TableStructurePrediction] = []
+
+        for page in pages:
+            assert page._backend is not None
+            if not page._backend.is_valid():
+                existing_prediction = page.predictions.tablestructure or TableStructurePrediction()
+                page.predictions.tablestructure = existing_prediction
+                predictions.append(existing_prediction)
+                continue
+
+            with TimeRecorder(conv_res, "table_structure"):
+                assert page.predictions.layout is not None
+                assert page.size is not None
+
+                table_prediction = TableStructurePrediction()
+                page.predictions.tablestructure = table_prediction
+
+                in_tables = [
+                    (
+                        cluster,
+                        [
+                            round(cluster.bbox.l) * self.scale,
+                            round(cluster.bbox.t) * self.scale,
+                            round(cluster.bbox.r) * self.scale,
+                            round(cluster.bbox.b) * self.scale,
+                        ],
+                    )
+                    for cluster in page.predictions.layout.clusters
+                    if cluster.label in [DocItemLabel.TABLE, DocItemLabel.DOCUMENT_INDEX]
+                ]
+                if not in_tables:
+                    predictions.append(table_prediction)
+                    continue
+
+                page_input: dict = {
+                    "width": page.size.width * self.scale,
+                    "height": page.size.height * self.scale,
+                    "image": numpy.asarray(page.get_image(scale=self.scale)),
+                }
+
+                for table_cluster, tbl_box in in_tables:
+                    page_input["tokens"] = get_tokens(page, table_cluster, self.scale)
+
+                    table_image = page_input["image"][
+                        round(tbl_box[1]) : round(tbl_box[3]),
+                        round(tbl_box[0]) : round(tbl_box[2]),
+                    ]
+
+                    table = self.pipeline([table_image])[0]
+
+                    docling_cells = []
+
+                    for cell_id, cell in enumerate(table.cells):
+                        docling_cell_bbox: dict = {
+                            "l": (cell.bbox.l + tbl_box[0]) / self.scale,
+                            "t": (cell.bbox.t + tbl_box[1]) / self.scale,
+                            "r": (cell.bbox.r + tbl_box[0]) / self.scale,
+                            "b": (cell.bbox.b + tbl_box[1]) / self.scale,
+                            "token": "",
+                        }
+
+                        docling_cell: dict = {
+                            "cell_id": cell_id,
+                            "bbox": docling_cell_bbox,
+                            "row_span": cell.row_span,
+                            "col_span": cell.col_span,
+                            "start_row_offset_idx": cell.row,
+                            "end_row_offset_idx": cell.row + cell.row_span,
+                            "start_col_offset_idx": cell.col,
+                            "end_col_offset_idx": cell.col + cell.col_span,
+                            "indentation_level": 0,
+                            "text_cell_bboxes": [docling_cell_bbox],
+                            "column_header": False,
+                            "row_header": False,
+                            "row_section": False,
+                        }
+
+                        bbox = BoundingBox.model_validate(docling_cell["bbox"])
+
+                        text_piece = page._backend.get_text_in_rect(bbox) if page._backend else ""
+                        docling_cell["bbox"]["token"] = text_piece
+
+                        tc = TableCell.model_validate(docling_cell)
+                        docling_cells.append(tc)
+
+                    docling_table = Table(
+                        otsl_seq=[],
+                        table_cells=docling_cells,
+                        num_rows=table.num_rows,
+                        num_cols=table.num_cols,
+                        id=table_cluster.id,
+                        page_no=page.page_no,
+                        cluster=table_cluster,
+                        label=table_cluster.label,
+                    )
+
+                    page.predictions.tablestructure.table_map[table_cluster.id] = docling_table
+
+                if settings.debug.visualize_tables:
+                    self.draw_table_and_cells(
+                        conv_res,
+                        page,
+                        page.predictions.tablestructure.table_map.values(),
+                    )
+
+                predictions.append(table_prediction)
+
+        return predictions
+
+    def draw_table_and_cells(
+        self,
+        conv_res: ConversionResult,
+        page: Page,
+        tbl_list: Iterable[Table],
+        show: bool = False,
+    ):
+        assert page._backend is not None
+        assert page.size is not None
+
+        image = page._backend.get_page_image()  # make new image to avoid drawing on the saved ones
+
+        scale_x = image.width / page.size.width
+        scale_y = image.height / page.size.height
+
+        draw = ImageDraw.Draw(image)
+
+        for table_element in tbl_list:
+            x0, y0, x1, y1 = table_element.cluster.bbox.as_tuple()
+            y0 *= scale_y
+            y1 *= scale_y
+            x0 *= scale_x
+            x1 *= scale_x
+
+            draw.rectangle([(x0, y0), (x1, y1)], outline="red")
+
+            for cell in table_element.cluster.cells:
+                x0, y0, x1, y1 = cell.rect.to_bounding_box().as_tuple()
+                x0 *= scale_x
+                x1 *= scale_x
+                y0 *= scale_y
+                y1 *= scale_y
+
+                draw.rectangle([(x0, y0), (x1, y1)], outline="green")
+
+            for tc in table_element.table_cells:
+                if tc.bbox is not None:
+                    x0, y0, x1, y1 = tc.bbox.as_tuple()
+                    x0 *= scale_x
+                    x1 *= scale_x
+                    y0 *= scale_y
+                    y1 *= scale_y
+
+                    if tc.column_header:
+                        width = 3
+                    else:
+                        width = 1
+                    draw.rectangle([(x0, y0), (x1, y1)], outline="blue", width=width)
+                    draw.text(
+                        (x0 + 3, y0 + 3),
+                        text=f"{tc.start_row_offset_idx}, {tc.start_col_offset_idx}",
+                        fill="black",
+                    )
+        if show:
+            image.show()
+        else:
+            out_path: Path = (
+                Path(settings.debug.debug_output_path) / f"debug_{conv_res.input.file.stem}"
+            )
+            out_path.mkdir(parents=True, exist_ok=True)
+
+            out_file = out_path / f"table_struct_page_{page.page_no:05}.png"
+            image.save(str(out_file), format="png")
+
+
+# Plugin factory
+def table_structure_engines():
+    return {"table_structure_engines": [CustomDoclingTableStructureModel]}

cells2table-0.2.0/src/cells2table/models/PaddlePaddle/__init__.py

@@ -0,0 +1,15 @@
+from .cell_detection import (
+    PaddlePaddleCellDetection,
+    PaddlePaddleWiredCellDetection,
+    PaddlePaddleWirelessCellDetection,
+)
+from .pipeline import PaddlePaddleTablePipeline
+from .table_classification import PaddlePaddleTableClassification
+
+__all__ = [
+    PaddlePaddleCellDetection,
+    PaddlePaddleWiredCellDetection,
+    PaddlePaddleWirelessCellDetection,
+    PaddlePaddleTablePipeline,
+    PaddlePaddleTableClassification,
+]

cells2table-0.2.0/src/cells2table/models/PaddlePaddle/cell_detection.py

@@ -0,0 +1,91 @@
+import logging
+from typing import Iterable, Iterator, Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..utils.download import DownloadOptions, DownloadPlatform
+from ..utils.runtimes import OnnxModel
+from ..utils.tasks import DetectionModel, DetectionResult
+
+HF_REPO_ID = "jspast/paddlepaddle-table-models-onnx"
+CONFIDENCE_THRESHOLD = 0.5
+
+logger = logging.getLogger(__name__)
+
+
+class PaddlePaddleCellDetection(DetectionModel, OnnxModel):
+    """Table cell detection model from PaddlePaddle."""
+
+    @property
+    def input_shape(self):
+        return self.session.get_inputs()[1].shape[2:]  # assuming NCHW
+
+    def __call__(self, input: Iterable[NDArray[np.uint8]]) -> list[Iterator[DetectionResult]]:
+        logger.debug("Started preprocessing")
+
+        original_shapes = []
+        scale_factors = []
+        for img in input:
+            original_shape = img.shape[:2]
+            original_shapes.append(original_shape)
+            scale_factors.append(
+                tuple(original_shape[i] / self.input_shape[i] for i in range(0, 2))
+            )
+
+        imgs = self.preprocess(input)
+
+        input_dict = dict(zip(self.input_names, [original_shapes, imgs, scale_factors]))
+
+        logger.debug("Done preprocessing")
+        logger.debug("Started running the model")
+
+        output = self.session.run(self.output_names, input_dict)
+
+        logger.debug("Done running the model")
+        logger.debug("Started postprocessing")
+
+        result = self.postprocess(output, scale_factors)  # type: ignore
+
+        logger.debug("Done postprocessing")
+
+        return result
+
+    def postprocess(
+        self,
+        pred: NDArray,
+        scale_factors: Sequence[tuple[int, int]],
+    ) -> list[Iterator[DetectionResult]]:
+        last_cell_idx = 0
+        batch_size = len(pred[1])
+        generators = []
+        cells = pred[0]
+
+        for i in range(batch_size):
+            c = cells[last_cell_idx : last_cell_idx + pred[1][i]]
+            c = c[c[:, 1] > CONFIDENCE_THRESHOLD]
+
+            last_cell_idx += pred[1][i]
+
+            if c.size:
+                sx, sy = scale_factors[i]
+                scores = c[:, 0]
+                boxes = c[:, 2:]
+                boxes[:, [0, 2]] *= sy
+                boxes[:, [1, 3]] *= sx
+
+            generators.append((DetectionResult(box, score) for box, score in zip(boxes, scores)))
+
+        return generators
+
+
+class PaddlePaddleWiredCellDetection(PaddlePaddleCellDetection):
+    download_options = DownloadOptions(
+        DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "wired_table_cell_det.onnx"
+    )
+
+
+class PaddlePaddleWirelessCellDetection(PaddlePaddleCellDetection):
+    download_options = DownloadOptions(
+        DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "wireless_table_cell_det.onnx"
+    )

cells2table-0.2.0/src/cells2table/models/PaddlePaddle/pipeline.py

@@ -0,0 +1,59 @@
+import logging
+from pathlib import Path
+from typing import Iterable, Optional
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ...utils.table import Table
+from .cell_detection import PaddlePaddleWiredCellDetection, PaddlePaddleWirelessCellDetection
+from .table_classification import PaddlePaddleTableClassification
+
+logger = logging.getLogger(__name__)
+
+
+class PaddlePaddleTablePipeline:
+    """A table pipeline combining PaddlePaddle classification and detection models."""
+
+    def __init__(self, models_path: Optional[Path | str] = None):
+        cls_path, wired_path, wireless_path = None, None, None
+
+        if models_path is not None:
+            models_path = Path(models_path)
+            cls_path = models_path / PaddlePaddleTableClassification.download_options.model_path
+            wired_path = models_path / PaddlePaddleWiredCellDetection.download_options.model_path
+            wireless_path = models_path / PaddlePaddleWiredCellDetection.download_options.model_path
+
+        self.cls_predictor = PaddlePaddleTableClassification(cls_path)
+        self.wired_predictor = PaddlePaddleWiredCellDetection(wired_path)
+        self.wireless_predictor = PaddlePaddleWirelessCellDetection(wireless_path)
+
+    def __call__(self, input: Iterable[NDArray[np.uint8]]) -> list[Table]:
+        wired_images, wireless_images, output = [], [], []
+
+        cls_result = self.cls_predictor(input)
+
+        for i, (img, p) in enumerate(zip(input, cls_result)):
+            (wired_images if p.cls == "wired" else wireless_images).append(img)
+            logger.info("Image %d classified as %s", i, p.cls)
+
+        if len(wired_images):
+            wired_cells = self.wired_predictor(wired_images)
+
+        if len(wireless_images):
+            wireless_cells = self.wireless_predictor(wireless_images)
+
+        wired_idx = 0
+        wireless_idx = 0
+
+        for i in range(len(cls_result)):
+            if cls_result[i].cls == "wired":
+                cells_det = wired_cells[wired_idx]
+                wired_idx += 1
+            else:
+                cells_det = wireless_cells[wireless_idx]
+                wireless_idx += 1
+
+            output.append(Table.from_detections(cells_det))
+
+        return output

cells2table-0.2.0/src/cells2table/models/PaddlePaddle/table_classification.py

@@ -0,0 +1,43 @@
+import logging
+from typing import Iterable, Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..utils.download import DownloadOptions, DownloadPlatform
+from ..utils.runtimes import OnnxModel
+from ..utils.tasks import ClassificationModel, ClassificationResult
+
+HF_REPO_ID = "jspast/paddlepaddle-table-models-onnx"
+
+logger = logging.getLogger(__name__)
+
+
+class PaddlePaddleTableClassification(ClassificationModel, OnnxModel):
+    download_options = DownloadOptions(DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "table_cls.onnx")
+
+    def __call__(self, input: Iterable[NDArray[np.uint8]]) -> list[ClassificationResult]:
+        logger.debug("Started preprocessing")
+        input = self.preprocess(input)
+
+        input_dict = dict(zip(self.input_names, [input]))
+
+        logger.debug("Done preprocessing")
+        logger.debug("Started running the model")
+
+        output = self.session.run(self.output_names, input_dict)[0]
+
+        logger.debug("Done running the model")
+        logger.debug("Started postprocessing")
+
+        result = self.postprocess(output)  # type: ignore
+
+        logger.debug("Done postprocessing")
+
+        return result
+
+    @staticmethod
+    def postprocess(pred: Sequence[Sequence[float]]) -> list[ClassificationResult]:
+        return [
+            ClassificationResult({0: "wired", 1: "wireless"}[np.argmax(p)], max(p)) for p in pred
+        ]

cells2table-0.2.0/src/cells2table/models/__init__.py

@@ -0,0 +1,3 @@
+from .PaddlePaddle import PaddlePaddleTablePipeline as DefaultPipeline
+
+__all__ = [DefaultPipeline]

cells2table-0.2.0/src/cells2table/models/utils/download.py

@@ -0,0 +1,41 @@
+import logging
+from enum import Enum
+from pathlib import Path
+from typing import NamedTuple
+
+logger = logging.getLogger(__name__)
+
+
+class DownloadPlatform(Enum):
+    HUGGINGFACE = "huggingface"
+
+
+class DownloadOptions(NamedTuple):
+    platform: DownloadPlatform
+    repo_id: str
+    model_path: str
+
+
+def download(options: DownloadOptions) -> Path:
+    match options.platform:
+        case DownloadPlatform.HUGGINGFACE:
+            path = download_hf_model(options.repo_id)
+
+    return path / options.model_path
+
+
+def download_hf_model(repo_id: str) -> Path:
+    """Download a repository from Hugging Face and return its path."""
+
+    try:
+        from huggingface_hub import snapshot_download
+        from huggingface_hub.utils import disable_progress_bars
+    except ImportError:
+        raise ImportError("huggingface_hub is not installed. Unable to download the model.")
+
+    disable_progress_bars()
+
+    logger.info("Downloading HF repo %s", repo_id)
+    download_path = snapshot_download(repo_id=repo_id)
+
+    return Path(download_path)

cells2table-0.2.0/src/cells2table/models/utils/runtimes/__init__.py

@@ -0,0 +1,3 @@
+from .onnx import OnnxModel
+
+__all__ = [OnnxModel]

cells2table-0.2.0/src/cells2table/models/utils/runtimes/onnx.py

@@ -0,0 +1,58 @@
+from abc import ABC
+from pathlib import Path
+from typing import Iterable, Optional
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+from numpy.typing import NDArray
+
+from ..tasks.base import BaseModel
+
+
+class OnnxModel(BaseModel, ABC):
+    """Base interface for ONNX models."""
+
+    def __init__(self, model_path: Optional[Path | str] = None) -> None:
+        if model_path is None:
+           model_path = self.download()
+
+        providers_priority = [
+            "CUDAExecutionProvider",
+            "MIGraphXExecutionProvider",
+            "OpenVINOExecutionProvider",
+            "CPUExecutionProvider",
+        ]
+        available_providers = ort.get_available_providers()  # type: ignore
+
+        self.session = ort.InferenceSession(
+            model_path,
+            providers=[p for p in providers_priority if p in available_providers],
+        )
+
+        self.scale = 1 / 255.0
+        self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
+        self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
+
+    @property
+    def input_shape(self):
+        return self.session.get_inputs()[0].shape[2:]  # assuming NCHW
+
+    @property
+    def input_names(self):
+        return [v.name for v in self.session.get_inputs()]
+
+    @property
+    def output_names(self):
+        return [v.name for v in self.session.get_outputs()]
+
+    def preprocess(self, input: Iterable[NDArray[np.uint8]]) -> list[NDArray[np.uint8]]:
+        output = []
+
+        for img in input:
+            img = cv2.resize(img, dsize=self.input_shape, interpolation=cv2.INTER_LANCZOS4)
+            img = (img.astype(np.float32) * self.scale - self.mean) / self.std  # Normalize
+            img = img.transpose(2, 0, 1)  # HWC to CHW
+            output.append(img)
+
+        return output

cells2table-0.2.0/src/cells2table/models/utils/tasks/__init__.py

@@ -0,0 +1,5 @@
+from .base import BaseModel
+from .classification import ClassificationModel, ClassificationResult
+from .detection import DetectionModel, DetectionResult
+
+__all__ = [BaseModel, ClassificationModel, ClassificationResult, DetectionModel, DetectionResult]

cells2table-0.2.0/src/cells2table/models/utils/tasks/base.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Optional
+
+from ..download import DownloadOptions, download
+
+
+class BaseModel(ABC):
+    """Base interface for models of any type."""
+
+    download_options: Optional[DownloadOptions] = None
+
+    @abstractmethod
+    def __init__(self, model_path: Optional[Path | str] = None) -> None:
+        pass
+
+    @abstractmethod
+    def __call__(self, input: Any):
+        pass
+
+    @classmethod
+    def download(cls) -> Path:
+        if cls.download_options is not None:
+            return download(cls.download_options)
+        else:
+            raise NotImplementedError(
+                "Download is not implemented for this model. Please provide a path."
+            )

cells2table-0.2.0/src/cells2table/models/utils/tasks/classification.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+from typing import Any, NamedTuple
+
+from .base import BaseModel
+
+
+class ClassificationResult(NamedTuple):
+    """Result type for classification models."""
+
+    cls: str
+    confidence: float
+
+
+class ClassificationModel(BaseModel, ABC):
+    """Base interface for table classification models."""
+
+    @abstractmethod
+    def __call__(self, input: Any) -> list[ClassificationResult]:
+        pass

cells2table-0.2.0/src/cells2table/models/utils/tasks/detection.py

@@ -0,0 +1,21 @@
+from abc import ABC, abstractmethod
+from typing import Any, Iterator, NamedTuple
+
+import numpy as np
+
+from .base import BaseModel
+
+
+class DetectionResult(NamedTuple):
+    """Result type for a detection with no class."""
+
+    bbox: np.ndarray
+    confidence: float
+
+
+class DetectionModel(BaseModel, ABC):
+    """Base interface for detection models."""
+
+    @abstractmethod
+    def __call__(self, input: Any) -> list[Iterator[DetectionResult]]:
+        pass

cells2table-0.2.0/src/cells2table/utils/table.py

@@ -0,0 +1,140 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Iterable, Optional
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+from ..models.utils.tasks import DetectionResult
+
+
+@dataclass
+class BoundingBox:
+    l: float  # noqa: E741
+    t: float
+    r: float
+    b: float
+
+    @staticmethod
+    def from_array(bbox: ArrayLike[float]) -> BoundingBox:
+        return BoundingBox(l=bbox[0], t=bbox[1], r=bbox[2], b=bbox[3])
+
+    def as_array(self) -> ArrayLike[float]:
+        return np.array([self.l, self.t, self.r, self.b])
+
+
+@dataclass
+class Cell:
+    bbox: BoundingBox
+    row: int
+    col: int
+    row_span: int = 1
+    col_span: int = 1
+
+
+@dataclass
+class Table:
+    cells: list[Cell] = field(default_factory=list)
+    num_rows: int = 0
+    num_cols: int = 0
+
+    @staticmethod
+    def from_detections(cells_det: Iterable[DetectionResult], tolerance: float = 10) -> Table:
+        table = Table()
+
+        for cell_det in cells_det:
+            bbox = BoundingBox.from_array(cell_det.bbox)
+            cell = Cell(bbox=bbox, row=0, col=0)
+            table.cells.append(cell)
+
+        table.compute_rows_and_cols(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, cells: Optional[Iterable[Cell]] = None) -> list[Cell]:
+        if cells is None:
+            cells = self.cells
+
+        return sorted(self.cells, key=lambda cell: cell.bbox.t)
+
+    def sort_cells_by_cols(self, cells: Optional[Iterable[Cell]] = None) -> list[Cell]:
+        if cells is None:
+            cells = self.cells
+
+        return sorted(self.cells, key=lambda cell: cell.bbox.l)
+
+    def compute_rows(self, tolerance: float) -> None:
+        self.cells = 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.cells = 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

cells2table-0.2.0/src/cells2table/utils/visualize.py

@@ -0,0 +1,42 @@
+import cv2
+import numpy as np
+from numpy.typing import NDArray
+
+from .table import Table
+
+
+def visualize_table(
+    image: NDArray[np.uint8],
+    table: Table,
+    color=(0, 255, 0),
+    thickness=2,
+    window_name="Bounding Boxes",
+) -> None:
+    """Simple table visualization on top of the image.
+
+    image: np.ndarray (BGR image loaded with cv2)
+    """
+
+    img = image.copy()
+
+    for cell in table.cells:
+        cv2.rectangle(
+            img,
+            (round(cell.bbox.l), round(cell.bbox.t)),
+            (round(cell.bbox.r), round(cell.bbox.b)),
+            color,
+            thickness,
+        )
+        cv2.putText(
+            img,
+            f"{cell.row},{cell.col} : {cell.row_span},{cell.col_span}",
+            (round(cell.bbox.l), round(cell.bbox.t) + 10),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.5,
+            (128, 192, 0),
+            thickness,
+        )
+
+    cv2.imshow(window_name, img)
+    cv2.waitKey(0)
+    cv2.destroyAllWindows()