cells2table 0.2.2__tar.gz → 0.3.0__tar.gz

{cells2table-0.2.2 → cells2table-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cells2table
-Version: 0.2.2
+Version: 0.3.0
 Summary: Table image parsing with cell detection models
 Keywords: docling,plugin
 Author: jspast
@@ -12,7 +12,7 @@ Classifier: Programming Language :: Python :: 3
 Requires-Dist: numpy>=2.2.0,<3.0.0
 Requires-Dist: opencv-python>=4.11.0.86,<5.0.0.0
 Requires-Dist: docling>=2.66.0,<3.0.0 ; extra == 'docling'
-Requires-Dist: huggingface-hub>=0.36.0,<1.0.0 ; extra == 'huggingface'
+Requires-Dist: huggingface-hub>=0.36.0,<2.0.0 ; extra == 'huggingface'
 Requires-Dist: onnxruntime>=1.23.2,<2.0.0 ; extra == 'onnx-cpu'
 Requires-Dist: onnxruntime-gpu>=1.23.2,<2.0.0 ; extra == 'onnx-cuda'
 Requires-Dist: onnxruntime-openvino>=1.23.0,<2.0.0 ; extra == 'onnx-openvino'

cells2table-0.3.0/cells2table/cli.py

@@ -0,0 +1,64 @@
+import argparse
+import logging
+from pathlib import Path
+
+import cv2
+
+from cells2table.pipelines import DefaultPipeline
+from cells2table.utils.visualize import visualize_table
+
+logger = logging.getLogger(__name__)
+
+
+def download(local_dir: Path | str | None = None) -> None:
+    """Download default pipeline models."""
+
+    log_format = "%(asctime)s\t%(levelname)s\t%(name)s: %(message)s"
+    logging.basicConfig(level=logging.INFO, format=log_format)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--local-dir", type=Path, default=None, help="Path to download models to")
+
+    args = parser.parse_args()
+
+    DefaultPipeline.download(local_dir=args.local_dir)
+
+
+def main() -> None:
+    """Basic CLI program for testing."""
+
+    log_format = "%(asctime)s\t%(levelname)s\t%(name)s: %(message)s"
+    logging.basicConfig(level=logging.INFO, format=log_format)
+
+    parser = argparse.ArgumentParser(description="Load an image from a given path using OpenCV")
+    parser.add_argument("image_path", type=Path, help="Path to the image file")
+    parser.add_argument("--models-path", type=Path, default=None, help="Path to downloaded models")
+
+    args = parser.parse_args()
+
+    if not args.image_path.exists():
+        raise FileNotFoundError(f"File does not exist: {args.image_path}")
+
+    image = cv2.imread(str(args.image_path))
+
+    if image is None:
+        raise ValueError(f"Failed to load image: {args.image_path}")
+
+    logger.info("Image loaded successfully from %s", args.image_path)
+    logger.debug(
+        "Image proprieties: width=%d, height=%d, channels=%d, datatype=%s",
+        image.shape[1],
+        image.shape[0],
+        image.shape[2],
+        str(image.dtype),
+    )
+
+    table_pipeline = DefaultPipeline(args.models_path)
+    tables = table_pipeline([image])
+
+    for table in tables:
+        visualize_table(image, table)  # type: ignore
+
+
+if __name__ == "__main__":
+    main()

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/datamodels/table.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import Iterable, Optional
+from typing import Iterable
 
-from ..models.tasks import DetectionResult
-from .bbox import BoundingBox
+from cells2table.datamodels import BoundingBox
+from cells2table.models.tasks import DetectionResult
 
 
 @dataclass
@@ -22,9 +22,9 @@ class Table:
     num_rows: int = 0
     num_cols: int = 0
 
-    @staticmethod
-    def from_detections(cells_det: Iterable[DetectionResult], tolerance: float = 10) -> Table:
-        table = Table()
+    @classmethod
+    def from_detections(cls, cells_det: Iterable[DetectionResult], tolerance: float = 10) -> Table:
+        table = cls()
 
         for cell_det in cells_det:
             bbox = BoundingBox.from_array(cell_det.bbox)
@@ -38,20 +38,14 @@ class Table:
         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
+    def sort_cells_by_rows(self) -> None:
+        self.cells = sorted(self.cells, key=lambda cell: cell.bbox.t)
 
-        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 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.cells = self.sort_cells_by_rows()
+        self.sort_cells_by_rows()
 
         row_y = None
         row_num = 0
@@ -87,7 +81,7 @@ class Table:
         self.num_rows = row_num + 1
 
     def compute_cols(self, tolerance: float) -> None:
-        self.cells = self.sort_cells_by_cols()
+        self.sort_cells_by_cols()
 
         col_x = None
         col_num = 0

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/docling.py

@@ -1,7 +1,7 @@
 import copy
 from collections.abc import Iterable
 from pathlib import Path
-from typing import ClassVar, Literal, Optional, Sequence, Type
+from typing import ClassVar, Literal, Sequence, Type
 
 import numpy
 
@@ -21,7 +21,7 @@ try:
 except ImportError:
     raise ImportError("docling is not installed. Unable to initialize plugin.")
 
-from . import DefaultPipeline
+from cells2table.pipelines import DefaultPipeline
 
 
 def get_tokens(page: Page, table_cluster: Cluster, scale: float) -> list[str]:
@@ -65,7 +65,7 @@ class CustomDoclingTableStructureModel(BaseTableStructureModel):
     def __init__(
         self,
         enabled: bool,
-        artifacts_path: Optional[Path],
+        artifacts_path: Path | None,
         options: CustomDoclingTableStructureOptions,
         accelerator_options: AcceleratorOptions,
         **kwargs,

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/models/PaddlePaddle/cell_detection.py

@@ -4,9 +4,9 @@ from typing import Iterable, Iterator, Sequence
 import numpy as np
 from numpy.typing import NDArray
 
-from ..download import DownloadOptions, DownloadPlatform
-from ..runtimes.onnx import OnnxModel
-from ..tasks import DetectionModel, DetectionResult
+from cells2table.models.runtimes.onnx import OnnxModel
+from cells2table.models.tasks import DetectionModel, DetectionResult
+from cells2table.utils.download import DownloadOptions, DownloadPlatform
 
 HF_REPO_ID = "jspast/paddlepaddle-table-models-onnx"
 CONFIDENCE_THRESHOLD = 0.5
@@ -58,22 +58,24 @@ class PaddlePaddleCellDetectionModel(DetectionModel, OnnxModel):
         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]]
+        generators = []
+
+        for i, count in enumerate(pred[1]):
+            c = cells[last_cell_idx : last_cell_idx + count]
             c = c[c[:, 1] > CONFIDENCE_THRESHOLD]
+            last_cell_idx += count
 
-            last_cell_idx += pred[1][i]
+            if not c.size:
+                generators.append(iter([]))
+                continue
 
-            if c.size:
-                sx, sy = scale_factors[i]
-                scores = c[:, 0]
-                boxes = c[:, 2:]
-                boxes[:, [0, 2]] *= sy
-                boxes[:, [1, 3]] *= sx
+            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)))
 
@@ -82,13 +84,23 @@ class PaddlePaddleCellDetectionModel(DetectionModel, OnnxModel):
 
 class PaddlePaddleWiredCellDetectionModel(PaddlePaddleCellDetectionModel):
     classes = ["wired"]
-    download_options = DownloadOptions(
-        DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "wired_table_cell_det.onnx"
-    )
+
+    @classmethod
+    def get_onnx_path(cls) -> str:
+        return "wired_table_cell_det.onnx"
+
+    @classmethod
+    def get_download_options(cls) -> DownloadOptions:
+        return DownloadOptions(DownloadPlatform.HUGGINGFACE, HF_REPO_ID, [cls.get_onnx_path()])
 
 
 class PaddlePaddleWirelessCellDetectionModel(PaddlePaddleCellDetectionModel):
     classes = ["wireless"]
-    download_options = DownloadOptions(
-        DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "wireless_table_cell_det.onnx"
-    )
+
+    @classmethod
+    def get_onnx_path(cls) -> str:
+        return "wireless_table_cell_det.onnx"
+
+    @classmethod
+    def get_download_options(cls) -> DownloadOptions:
+        return DownloadOptions(DownloadPlatform.HUGGINGFACE, HF_REPO_ID, [cls.get_onnx_path()])

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/models/PaddlePaddle/table_classification.py

@@ -4,9 +4,9 @@ from typing import Iterable, Sequence
 import numpy as np
 from numpy.typing import NDArray
 
-from ..download import DownloadOptions, DownloadPlatform
-from ..runtimes.onnx import OnnxModel
-from ..tasks import ClassificationModel, ClassificationResult
+from cells2table.models.runtimes.onnx import OnnxModel
+from cells2table.models.tasks import ClassificationModel, ClassificationResult
+from cells2table.utils.download import DownloadOptions, DownloadPlatform
 
 HF_REPO_ID = "jspast/paddlepaddle-table-models-onnx"
 
@@ -15,7 +15,14 @@ logger = logging.getLogger(__name__)
 
 class PaddlePaddleTableClassificationModel(ClassificationModel, OnnxModel):
     classes = ["wired", "wireless"]
-    download_options = DownloadOptions(DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "table_cls.onnx")
+
+    @classmethod
+    def get_onnx_path(cls) -> str:
+        return "table_cls.onnx"
+
+    @classmethod
+    def get_download_options(cls) -> DownloadOptions:
+        return DownloadOptions(DownloadPlatform.HUGGINGFACE, HF_REPO_ID, [cls.get_onnx_path()])
 
     def __call__(self, input: Iterable[NDArray[np.uint8]]) -> list[ClassificationResult]:
         logger.debug("Started preprocessing")

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/models/runtimes/onnx.py

@@ -1,13 +1,13 @@
-from abc import ABC
+from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Iterable, Optional
+from typing import Iterable
 
 import cv2
 import numpy as np
 import onnxruntime as ort
 from numpy.typing import NDArray
 
-from ..tasks.base import BaseModel
+from cells2table.models.tasks.base import BaseModel
 
 
 class OnnxModel(BaseModel, ABC):
@@ -17,7 +17,12 @@ class OnnxModel(BaseModel, ABC):
     mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
     std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
 
-    def __init__(self, model_path: Optional[Path | str] = None) -> None:
+    @classmethod
+    @abstractmethod
+    def get_onnx_path(self) -> str:
+        pass
+
+    def __init__(self, model_path: Path | str | None = None) -> None:
         self.model_path = self.download() if model_path is None else Path(model_path)
 
         providers_priority = [
@@ -26,10 +31,10 @@ class OnnxModel(BaseModel, ABC):
             "OpenVINOExecutionProvider",
             "CPUExecutionProvider",
         ]
-        available_providers = ort.get_available_providers()  # type: ignore
+        available_providers = ort.get_available_providers()
 
         self.session = ort.InferenceSession(
-            self.model_path,
+            self.model_path / self.get_onnx_path(),
             providers=[p for p in providers_priority if p in available_providers],
         )
 

cells2table-0.3.0/cells2table/models/tasks/base.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+from cells2table.utils.download import DownloadOptions
+
+
+class BaseModel(ABC):
+    """Base interface for models of any type."""
+
+    model_path: Path
+
+    @abstractmethod
+    def __init__(self, model_path: Path | str | None = None) -> None:
+        pass
+
+    @abstractmethod
+    def __call__(self, input: Any):
+        pass
+
+    @classmethod
+    @abstractmethod
+    def get_download_options(cls) -> DownloadOptions:
+        pass
+
+    @classmethod
+    def download(cls, *, local_dir: Path | str | None = None) -> Path:
+        return cls.get_download_options().download(local_dir=local_dir)

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/models/tasks/classification.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Any, NamedTuple
 
-from .base import BaseModel
+from cells2table.models.tasks.base import BaseModel
 
 
 class ClassificationResult(NamedTuple):

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/models/tasks/detection.py

@@ -3,7 +3,7 @@ from typing import Any, Iterator, NamedTuple
 
 import numpy as np
 
-from .base import BaseModel
+from cells2table.models.tasks.base import BaseModel
 
 
 class DetectionResult(NamedTuple):

cells2table-0.3.0/cells2table/pipelines/__init__.py

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

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/pipelines/base.py

@@ -1,15 +1,20 @@
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 
 class BasePipeline(ABC):
     """Base interface for pipelines of any type."""
 
     @abstractmethod
-    def __init__(self, models_path: Optional[Path | str] = None) -> None:
+    def __init__(self, models_path: Path | str | None = None) -> None:
         pass
 
     @abstractmethod
     def __call__(self, input: Any):
         pass
+
+    @classmethod
+    @abstractmethod
+    def download(cls, *, local_dir: Path | str | None = None) -> Path:
+        pass

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/pipelines/classification_detection.py

@@ -1,11 +1,11 @@
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, Iterable, Optional
+from typing import Any, Iterable
 
-from ..datamodels import Table
-from ..models.tasks import ClassificationModel, DetectionModel
-from .base import BasePipeline
+from cells2table.datamodels import Table
+from cells2table.models.tasks import ClassificationModel, DetectionModel
+from cells2table.pipelines.base import BasePipeline
 
 logger = logging.getLogger(__name__)
 
@@ -17,7 +17,7 @@ class ClassificationDetectionPipeline(BasePipeline, ABC):
     detection_models: list[DetectionModel]
 
     @abstractmethod
-    def __init__(self, models_path: Optional[Path | str] = None) -> None:
+    def __init__(self, models_path: Path | str | None = None) -> None:
         """Initialize the models."""
 
         pass

cells2table-0.3.0/cells2table/pipelines/paddlepaddle.py

@@ -0,0 +1,30 @@
+from pathlib import Path
+
+from cells2table.models.PaddlePaddle import (
+    PaddlePaddleTableClassificationModel,
+    PaddlePaddleWiredCellDetectionModel,
+    PaddlePaddleWirelessCellDetectionModel,
+)
+from cells2table.pipelines.classification_detection import ClassificationDetectionPipeline
+
+
+class PaddlePaddleTablePipeline(ClassificationDetectionPipeline):
+    """A table pipeline combining PaddlePaddle classification and detection models."""
+
+    def __init__(self, models_path: Path | str | None = None) -> None:
+        """Initialize models from the provided path or download them."""
+
+        models_path = self.download() if models_path is None else Path(models_path)
+
+        self.classification_model = PaddlePaddleTableClassificationModel(models_path)
+        self.detection_models = [
+            PaddlePaddleWiredCellDetectionModel(models_path),
+            PaddlePaddleWirelessCellDetectionModel(models_path),
+        ]
+
+    @classmethod
+    def download(cls, *, local_dir: Path | str | None = None) -> Path:
+        PaddlePaddleTableClassificationModel.download(local_dir=local_dir)
+        PaddlePaddleWiredCellDetectionModel.download(local_dir=local_dir)
+        path = PaddlePaddleWirelessCellDetectionModel.download(local_dir=local_dir)
+        return path

{cells2table-0.2.2/src/cells2table/models → cells2table-0.3.0/cells2table/utils}/download.py

@@ -13,18 +13,22 @@ class DownloadPlatform(Enum):
 class DownloadOptions(NamedTuple):
     platform: DownloadPlatform
     repo_id: str
-    model_path: str
+    files: list[str] | None = None
 
+    def download(self, *, local_dir: Path | str | None = None) -> Path:
+        match self.platform:
+            case DownloadPlatform.HUGGINGFACE:
+                path = download_hf_model(self.repo_id, files=self.files, local_dir=local_dir)
 
-def download(options: DownloadOptions) -> Path:
-    match options.platform:
-        case DownloadPlatform.HUGGINGFACE:
-            path = download_hf_model(options.repo_id)
+        return path
 
-    return path / options.model_path
 
-
-def download_hf_model(repo_id: str) -> Path:
+def download_hf_model(
+    repo_id: str,
+    *,
+    files: list[str] | None = None,
+    local_dir: Path | str | None = None,
+) -> Path:
     """Download a repository from Hugging Face and return its path."""
 
     try:
@@ -36,6 +40,6 @@ def download_hf_model(repo_id: str) -> Path:
     disable_progress_bars()
 
     logger.info("Downloading HF repo %s", repo_id)
-    download_path = snapshot_download(repo_id=repo_id)
+    download_path = snapshot_download(repo_id=repo_id, allow_patterns=files, local_dir=local_dir)
 
     return Path(download_path)

{cells2table-0.2.2/src → cells2table-0.3.0}/cells2table/utils/visualize.py

@@ -2,7 +2,7 @@ import cv2
 import numpy as np
 from numpy.typing import NDArray
 
-from ..datamodels import Table
+from cells2table.datamodels import Table
 
 
 def visualize_table(

{cells2table-0.2.2 → cells2table-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "cells2table"
-version = "0.2.2"
+version = "0.3.0"
 description = "Table image parsing with cell detection models"
 readme = "README.md"
 requires-python = ">=3.12,<4.0"
@@ -24,7 +24,8 @@ Homepage = "https://github.com/jspast/cells2table"
 Issues = "https://github.com/jspast/cells2table/issues"
 
 [project.scripts]
-cells2table = "cli:main"
+cells2table = "cells2table.cli:main"
+cells2table-download = "cells2table.cli:download"
 
 [project.entry-points.docling]
 cells2table = "cells2table.docling"
@@ -34,7 +35,7 @@ docling = [
   "docling (>=2.66.0,<3.0.0)",
 ]
 huggingface = [
-  "huggingface-hub (>=0.36.0,<1.0.0)",
+  "huggingface-hub (>=0.36.0,<2.0.0)",
 ]
 onnx_cpu = [
   "onnxruntime (>=1.23.2,<2.0.0)",
@@ -49,12 +50,12 @@ onnx_openvino = [
 [dependency-groups]
 dev = [
   "cells2table",
-  "pytest (>=7.0,<8.0)",
+  "pytest (>=7.0,<10.0)",
   "ruff (>=0.5.0,<1.0.0)",
 ]
 
 [build-system]
-requires = ["uv_build>=0.9.17,<0.10.0"]
+requires = ["uv_build>=0.9.17,<1.0.0"]
 build-backend = "uv_build"
 
 [tool.ruff]
@@ -69,3 +70,7 @@ conflicts = [
     { extra = "onnx_openvino" }
   ]
 ]
+
+[tool.uv.build-backend]
+module-name = "cells2table"
+module-root = ""

cells2table-0.2.2/src/cells2table/__init__.py

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

cells2table-0.2.2/src/cells2table/models/tasks/base.py

@@ -1,29 +0,0 @@
-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."""
-
-    model_path: Path
-    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.2/src/cells2table/pipelines/__init__.py

@@ -1,3 +0,0 @@
-from .paddlepaddle import PaddlePaddleTablePipeline
-
-__all__ = [PaddlePaddleTablePipeline]

cells2table-0.2.2/src/cells2table/pipelines/paddlepaddle.py

@@ -1,45 +0,0 @@
-import logging
-from pathlib import Path
-from typing import Optional
-
-from ..models.PaddlePaddle import (
-    PaddlePaddleTableClassificationModel,
-    PaddlePaddleWiredCellDetectionModel,
-    PaddlePaddleWirelessCellDetectionModel,
-)
-from .classification_detection import ClassificationDetectionPipeline
-
-logger = logging.getLogger(__name__)
-
-
-class PaddlePaddleTablePipeline(ClassificationDetectionPipeline):
-    """A table pipeline combining PaddlePaddle classification and detection models."""
-
-    def __init__(self, models_path: Optional[Path | str] = None) -> None:
-        """Initialize models from the provided path or download them.
-
-        As the models are all in the same repository, do the download only once.
-        """
-
-        self.detection_models = []
-        cls_path, wired_path, wireless_path = None, None, None
-
-        if models_path is not None:
-            models_path = Path(models_path)
-            cls_path = (
-                models_path / PaddlePaddleTableClassificationModel.download_options.model_path
-            )
-
-        self.classification_model = PaddlePaddleTableClassificationModel(cls_path)
-
-        wired_path = (
-            self.classification_model.model_path.parent
-            / PaddlePaddleWiredCellDetectionModel.download_options.model_path
-        )
-        wireless_path = (
-            self.classification_model.model_path.parent
-            / PaddlePaddleWirelessCellDetectionModel.download_options.model_path
-        )
-
-        self.detection_models.append(PaddlePaddleWiredCellDetectionModel(wired_path))
-        self.detection_models.append(PaddlePaddleWirelessCellDetectionModel(wireless_path))