cells2table 0.2.0__tar.gz → 0.2.2__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (32) hide show
  1. {cells2table-0.2.0 → cells2table-0.2.2}/PKG-INFO +23 -21
  2. {cells2table-0.2.0 → cells2table-0.2.2}/README.md +10 -9
  3. {cells2table-0.2.0 → cells2table-0.2.2}/pyproject.toml +30 -15
  4. cells2table-0.2.2/src/cells2table/__init__.py +3 -0
  5. cells2table-0.2.2/src/cells2table/datamodels/__init__.py +4 -0
  6. cells2table-0.2.2/src/cells2table/datamodels/bbox.py +21 -0
  7. {cells2table-0.2.0/src/cells2table/utils → cells2table-0.2.2/src/cells2table/datamodels}/table.py +2 -19
  8. {cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/docling.py +18 -19
  9. cells2table-0.2.2/src/cells2table/models/PaddlePaddle/__init__.py +13 -0
  10. {cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/models/PaddlePaddle/cell_detection.py +10 -7
  11. {cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/models/PaddlePaddle/table_classification.py +8 -9
  12. {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/runtimes/onnx.py +6 -7
  13. {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/tasks/base.py +1 -0
  14. {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/tasks/classification.py +2 -0
  15. {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/tasks/detection.py +2 -0
  16. cells2table-0.2.2/src/cells2table/pipelines/__init__.py +3 -0
  17. cells2table-0.2.2/src/cells2table/pipelines/base.py +15 -0
  18. cells2table-0.2.2/src/cells2table/pipelines/classification_detection.py +62 -0
  19. cells2table-0.2.2/src/cells2table/pipelines/paddlepaddle.py +45 -0
  20. cells2table-0.2.2/src/cells2table/py.typed +0 -0
  21. cells2table-0.2.2/src/cells2table/utils/__init__.py +0 -0
  22. {cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/utils/visualize.py +1 -1
  23. cells2table-0.2.0/src/cells2table/__init__.py +0 -3
  24. cells2table-0.2.0/src/cells2table/models/PaddlePaddle/__init__.py +0 -15
  25. cells2table-0.2.0/src/cells2table/models/PaddlePaddle/pipeline.py +0 -59
  26. cells2table-0.2.0/src/cells2table/models/__init__.py +0 -3
  27. cells2table-0.2.0/src/cells2table/models/utils/runtimes/__init__.py +0 -3
  28. {cells2table-0.2.0 → cells2table-0.2.2}/LICENSE +0 -0
  29. {cells2table-0.2.0/src/cells2table/utils → cells2table-0.2.2/src/cells2table/models}/__init__.py +0 -0
  30. {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/download.py +0 -0
  31. /cells2table-0.2.0/src/cells2table/py.typed → /cells2table-0.2.2/src/cells2table/models/runtimes/__init__.py +0 -0
  32. {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/tasks/__init__.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: cells2table
3
- Version: 0.2.0
3
+ Version: 0.2.2
4
4
  Summary: Table image parsing with cell detection models
5
5
  Keywords: docling,plugin
6
6
  Author: jspast
@@ -9,20 +9,21 @@ License-Expression: MIT
9
9
  License-File: LICENSE
10
10
  Classifier: Operating System :: OS Independent
11
11
  Classifier: Programming Language :: Python :: 3
12
- Requires-Dist: huggingface-hub>=0.36.0
13
- Requires-Dist: numpy>=2.2.0
14
- Requires-Dist: opencv-python>=4.11.0.86
15
- Requires-Dist: onnxruntime>=1.23.2 ; extra == 'cpu'
16
- Requires-Dist: onnxruntime-gpu>=1.23.2 ; extra == 'cuda'
17
- Requires-Dist: docling>=2.66.0 ; extra == 'docling'
18
- Requires-Dist: onnxruntime-openvino>=1.23.0 ; extra == 'openvino'
19
- Requires-Python: >=3.12
12
+ Requires-Dist: numpy>=2.2.0,<3.0.0
13
+ Requires-Dist: opencv-python>=4.11.0.86,<5.0.0.0
14
+ Requires-Dist: docling>=2.66.0,<3.0.0 ; extra == 'docling'
15
+ Requires-Dist: huggingface-hub>=0.36.0,<1.0.0 ; extra == 'huggingface'
16
+ Requires-Dist: onnxruntime>=1.23.2,<2.0.0 ; extra == 'onnx-cpu'
17
+ Requires-Dist: onnxruntime-gpu>=1.23.2,<2.0.0 ; extra == 'onnx-cuda'
18
+ Requires-Dist: onnxruntime-openvino>=1.23.0,<2.0.0 ; extra == 'onnx-openvino'
19
+ Requires-Python: >=3.12, <4.0
20
20
  Project-URL: Homepage, https://github.com/jspast/cells2table
21
21
  Project-URL: Issues, https://github.com/jspast/cells2table/issues
22
- Provides-Extra: cpu
23
- Provides-Extra: cuda
24
22
  Provides-Extra: docling
25
- Provides-Extra: openvino
23
+ Provides-Extra: huggingface
24
+ Provides-Extra: onnx-cpu
25
+ Provides-Extra: onnx-cuda
26
+ Provides-Extra: onnx-openvino
26
27
  Description-Content-Type: text/markdown
27
28
 
28
29
  # cells2table
@@ -31,29 +32,30 @@ Parsing tables in document images with cell detection models
31
32
 
32
33
  ## Implemented pipelines
33
34
 
34
- ### PaddlePaddle models
35
+ ### PaddlePaddle
35
36
 
36
37
  - Classification model (wired / wireless)
37
38
  - Cell detection model with different weights for each class
38
39
 
39
- Using [ONNX weights](https://huggingface.co/jspast/paddlepaddle-table-models-onnx) (downloaded automatically on first use with `huggingface_hub`)
40
+ Uses ONNX weights downloaded automatically from [Hugging Face](https://huggingface.co/jspast/paddlepaddle-table-models-onnx) on first use.
40
41
 
41
42
  ## Instalation
42
43
 
43
44
  With [uv](https://docs.astral.sh/uv/), add to your project with:
44
45
 
45
46
  ```sh
46
- uv add git+https://github.com/jspast/cells2table
47
+ uv add cells2table
47
48
  ```
48
49
 
49
50
  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.
50
51
 
51
- | Optional | Description |
52
- | ---------- | ----------------------- |
53
- | `cuda` | For NVIDIA GPUs |
54
- | `openvino` | For Intel GPUs and CPUs |
55
- | `cpu` | Default CPU runtime |
56
- | `docling` | For docling usage |
52
+ | Optional | Description |
53
+ | --------------- | ----------------------- |
54
+ | `docling` | For docling usage |
55
+ | `huggingface` | For downloading models |
56
+ | `onnx_cuda` | For NVIDIA GPUs |
57
+ | `onnx_openvino` | For Intel GPUs and CPUs |
58
+ | `onnx_cpu` | Default CPU runtime |
57
59
 
58
60
  ## Usage
59
61
 
@@ -4,29 +4,30 @@ Parsing tables in document images with cell detection models
4
4
 
5
5
  ## Implemented pipelines
6
6
 
7
- ### PaddlePaddle models
7
+ ### PaddlePaddle
8
8
 
9
9
  - Classification model (wired / wireless)
10
10
  - Cell detection model with different weights for each class
11
11
 
12
- Using [ONNX weights](https://huggingface.co/jspast/paddlepaddle-table-models-onnx) (downloaded automatically on first use with `huggingface_hub`)
12
+ Uses ONNX weights downloaded automatically from [Hugging Face](https://huggingface.co/jspast/paddlepaddle-table-models-onnx) on first use.
13
13
 
14
14
  ## Instalation
15
15
 
16
16
  With [uv](https://docs.astral.sh/uv/), add to your project with:
17
17
 
18
18
  ```sh
19
- uv add git+https://github.com/jspast/cells2table
19
+ uv add cells2table
20
20
  ```
21
21
 
22
22
  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.
23
23
 
24
- | Optional | Description |
25
- | ---------- | ----------------------- |
26
- | `cuda` | For NVIDIA GPUs |
27
- | `openvino` | For Intel GPUs and CPUs |
28
- | `cpu` | Default CPU runtime |
29
- | `docling` | For docling usage |
24
+ | Optional | Description |
25
+ | --------------- | ----------------------- |
26
+ | `docling` | For docling usage |
27
+ | `huggingface` | For downloading models |
28
+ | `onnx_cuda` | For NVIDIA GPUs |
29
+ | `onnx_openvino` | For Intel GPUs and CPUs |
30
+ | `onnx_cpu` | Default CPU runtime |
30
31
 
31
32
  ## Usage
32
33
 
@@ -1,9 +1,9 @@
1
1
  [project]
2
2
  name = "cells2table"
3
- version = "0.2.0"
3
+ version = "0.2.2"
4
4
  description = "Table image parsing with cell detection models"
5
5
  readme = "README.md"
6
- requires-python = ">=3.12"
6
+ requires-python = ">=3.12,<4.0"
7
7
  license = "MIT"
8
8
  license-files = ["LICENSE"]
9
9
  authors = [
@@ -15,9 +15,8 @@ classifiers = [
15
15
  "Programming Language :: Python :: 3",
16
16
  ]
17
17
  dependencies = [
18
- "huggingface-hub>=0.36.0",
19
- "numpy>=2.2.0",
20
- "opencv-python>=4.11.0.86",
18
+ "numpy (>=2.2.0,<3.0.0)",
19
+ "opencv-python (>=4.11.0.86,<5.0.0.0)",
21
20
  ]
22
21
 
23
22
  [project.urls]
@@ -25,23 +24,33 @@ Homepage = "https://github.com/jspast/cells2table"
25
24
  Issues = "https://github.com/jspast/cells2table/issues"
26
25
 
27
26
  [project.scripts]
28
- cells2table = "cells2table.cli:main"
27
+ cells2table = "cli:main"
29
28
 
30
29
  [project.entry-points.docling]
31
30
  cells2table = "cells2table.docling"
32
31
 
33
32
  [project.optional-dependencies]
34
- cpu = [
35
- "onnxruntime>=1.23.2",
33
+ docling = [
34
+ "docling (>=2.66.0,<3.0.0)",
36
35
  ]
37
- cuda = [
38
- "onnxruntime-gpu>=1.23.2",
36
+ huggingface = [
37
+ "huggingface-hub (>=0.36.0,<1.0.0)",
39
38
  ]
40
- docling = [
41
- "docling>=2.66.0",
39
+ onnx_cpu = [
40
+ "onnxruntime (>=1.23.2,<2.0.0)",
41
+ ]
42
+ onnx_cuda = [
43
+ "onnxruntime-gpu (>=1.23.2,<2.0.0)",
42
44
  ]
43
- openvino = [
44
- "onnxruntime-openvino>=1.23.0",
45
+ onnx_openvino = [
46
+ "onnxruntime-openvino (>=1.23.0,<2.0.0)",
47
+ ]
48
+
49
+ [dependency-groups]
50
+ dev = [
51
+ "cells2table",
52
+ "pytest (>=7.0,<8.0)",
53
+ "ruff (>=0.5.0,<1.0.0)",
45
54
  ]
46
55
 
47
56
  [build-system]
@@ -53,4 +62,10 @@ line-length = 100
53
62
 
54
63
  [tool.uv]
55
64
  package = true
56
- conflicts = [[{ extra = "cpu" }, { extra = "cuda" }, { extra = "openvino" }]]
65
+ conflicts = [
66
+ [
67
+ { extra = "onnx_cpu" },
68
+ { extra = "onnx_cuda" },
69
+ { extra = "onnx_openvino" }
70
+ ]
71
+ ]
@@ -0,0 +1,3 @@
1
+ from .pipelines import PaddlePaddleTablePipeline as DefaultPipeline
2
+
3
+ __all__ = [DefaultPipeline]
@@ -0,0 +1,4 @@
1
+ from .bbox import BoundingBox
2
+ from .table import Cell, Table
3
+
4
+ __all__ = [BoundingBox, Cell, Table]
@@ -0,0 +1,21 @@
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import dataclass
4
+
5
+ import numpy as np
6
+ from numpy.typing import ArrayLike, NDArray
7
+
8
+
9
+ @dataclass
10
+ class BoundingBox:
11
+ l: float # noqa: E741
12
+ t: float
13
+ r: float
14
+ b: float
15
+
16
+ @staticmethod
17
+ def from_array(bbox: ArrayLike[float]) -> BoundingBox:
18
+ return BoundingBox(l=bbox[0], t=bbox[1], r=bbox[2], b=bbox[3])
19
+
20
+ def as_array(self) -> NDArray[np.float32]:
21
+ return np.array([self.l, self.t, self.r, self.b], dtype=np.float32)
@@ -3,25 +3,8 @@ from __future__ import annotations
3
3
  from dataclasses import dataclass, field
4
4
  from typing import Iterable, Optional
5
5
 
6
- import numpy as np
7
- from numpy.typing import ArrayLike
8
-
9
- from ..models.utils.tasks import DetectionResult
10
-
11
-
12
- @dataclass
13
- class BoundingBox:
14
- l: float # noqa: E741
15
- t: float
16
- r: float
17
- b: float
18
-
19
- @staticmethod
20
- def from_array(bbox: ArrayLike[float]) -> BoundingBox:
21
- return BoundingBox(l=bbox[0], t=bbox[1], r=bbox[2], b=bbox[3])
22
-
23
- def as_array(self) -> ArrayLike[float]:
24
- return np.array([self.l, self.t, self.r, self.b])
6
+ from ..models.tasks import DetectionResult
7
+ from .bbox import BoundingBox
25
8
 
26
9
 
27
10
  @dataclass
@@ -1,32 +1,30 @@
1
1
  import copy
2
- import logging
3
2
  from collections.abc import Iterable
4
3
  from pathlib import Path
5
4
  from typing import ClassVar, Literal, Optional, Sequence, Type
6
5
 
7
6
  import numpy
8
- from docling.datamodel.accelerator_options import AcceleratorOptions
9
- from docling.datamodel.base_models import Cluster, Page, Table, TableStructurePrediction
10
- from docling.datamodel.document import ConversionResult
11
- from docling.datamodel.pipeline_options import BaseTableStructureOptions
12
- from docling.datamodel.settings import settings
13
- from docling.models.base_table_model import BaseTableStructureModel
14
- from docling.utils.profiling import TimeRecorder
15
- from docling_core.types.doc.base import BoundingBox
16
- from docling_core.types.doc.document import TableCell
17
- from docling_core.types.doc.labels import DocItemLabel
18
- from docling_core.types.doc.page import BoundingRectangle, TextCellUnit
19
- from PIL import ImageDraw
20
7
 
21
- # from docling.utils.accelerator_utils import decide_device
22
- from .models import DefaultPipeline
23
-
24
- logger = logging.getLogger(__name__)
8
+ try:
9
+ from docling.datamodel.accelerator_options import AcceleratorOptions
10
+ from docling.datamodel.base_models import Cluster, Page, Table, TableStructurePrediction
11
+ from docling.datamodel.document import ConversionResult
12
+ from docling.datamodel.pipeline_options import BaseTableStructureOptions
13
+ from docling.datamodel.settings import settings
14
+ from docling.models.base_table_model import BaseTableStructureModel
15
+ from docling.utils.profiling import TimeRecorder
16
+ from docling_core.types.doc.base import BoundingBox
17
+ from docling_core.types.doc.document import TableCell
18
+ from docling_core.types.doc.labels import DocItemLabel
19
+ from docling_core.types.doc.page import BoundingRectangle, TextCellUnit
20
+ from PIL import ImageDraw
21
+ except ImportError:
22
+ raise ImportError("docling is not installed. Unable to initialize plugin.")
23
+
24
+ from . import DefaultPipeline
25
25
 
26
26
 
27
27
  def get_tokens(page: Page, table_cluster: Cluster, scale: float) -> list[str]:
28
- """Docling logic for token extraction."""
29
-
30
28
  # Check if word-level cells are available from backend:
31
29
  sp = page._backend.get_segmented_page() if page._backend else None
32
30
  if sp is not None:
@@ -70,6 +68,7 @@ class CustomDoclingTableStructureModel(BaseTableStructureModel):
70
68
  artifacts_path: Optional[Path],
71
69
  options: CustomDoclingTableStructureOptions,
72
70
  accelerator_options: AcceleratorOptions,
71
+ **kwargs,
73
72
  ):
74
73
  self.enabled = enabled
75
74
 
@@ -0,0 +1,13 @@
1
+ from .cell_detection import (
2
+ PaddlePaddleCellDetectionModel,
3
+ PaddlePaddleWiredCellDetectionModel,
4
+ PaddlePaddleWirelessCellDetectionModel,
5
+ )
6
+ from .table_classification import PaddlePaddleTableClassificationModel
7
+
8
+ __all__ = [
9
+ PaddlePaddleCellDetectionModel,
10
+ PaddlePaddleWiredCellDetectionModel,
11
+ PaddlePaddleWirelessCellDetectionModel,
12
+ PaddlePaddleTableClassificationModel,
13
+ ]
@@ -4,9 +4,9 @@ from typing import Iterable, Iterator, Sequence
4
4
  import numpy as np
5
5
  from numpy.typing import NDArray
6
6
 
7
- from ..utils.download import DownloadOptions, DownloadPlatform
8
- from ..utils.runtimes import OnnxModel
9
- from ..utils.tasks import DetectionModel, DetectionResult
7
+ from ..download import DownloadOptions, DownloadPlatform
8
+ from ..runtimes.onnx import OnnxModel
9
+ from ..tasks import DetectionModel, DetectionResult
10
10
 
11
11
  HF_REPO_ID = "jspast/paddlepaddle-table-models-onnx"
12
12
  CONFIDENCE_THRESHOLD = 0.5
@@ -14,7 +14,7 @@ CONFIDENCE_THRESHOLD = 0.5
14
14
  logger = logging.getLogger(__name__)
15
15
 
16
16
 
17
- class PaddlePaddleCellDetection(DetectionModel, OnnxModel):
17
+ class PaddlePaddleCellDetectionModel(DetectionModel, OnnxModel):
18
18
  """Table cell detection model from PaddlePaddle."""
19
19
 
20
20
  @property
@@ -51,8 +51,9 @@ class PaddlePaddleCellDetection(DetectionModel, OnnxModel):
51
51
 
52
52
  return result
53
53
 
54
+ @classmethod
54
55
  def postprocess(
55
- self,
56
+ cls,
56
57
  pred: NDArray,
57
58
  scale_factors: Sequence[tuple[int, int]],
58
59
  ) -> list[Iterator[DetectionResult]]:
@@ -79,13 +80,15 @@ class PaddlePaddleCellDetection(DetectionModel, OnnxModel):
79
80
  return generators
80
81
 
81
82
 
82
- class PaddlePaddleWiredCellDetection(PaddlePaddleCellDetection):
83
+ class PaddlePaddleWiredCellDetectionModel(PaddlePaddleCellDetectionModel):
84
+ classes = ["wired"]
83
85
  download_options = DownloadOptions(
84
86
  DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "wired_table_cell_det.onnx"
85
87
  )
86
88
 
87
89
 
88
- class PaddlePaddleWirelessCellDetection(PaddlePaddleCellDetection):
90
+ class PaddlePaddleWirelessCellDetectionModel(PaddlePaddleCellDetectionModel):
91
+ classes = ["wireless"]
89
92
  download_options = DownloadOptions(
90
93
  DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "wireless_table_cell_det.onnx"
91
94
  )
@@ -4,16 +4,17 @@ from typing import Iterable, Sequence
4
4
  import numpy as np
5
5
  from numpy.typing import NDArray
6
6
 
7
- from ..utils.download import DownloadOptions, DownloadPlatform
8
- from ..utils.runtimes import OnnxModel
9
- from ..utils.tasks import ClassificationModel, ClassificationResult
7
+ from ..download import DownloadOptions, DownloadPlatform
8
+ from ..runtimes.onnx import OnnxModel
9
+ from ..tasks import ClassificationModel, ClassificationResult
10
10
 
11
11
  HF_REPO_ID = "jspast/paddlepaddle-table-models-onnx"
12
12
 
13
13
  logger = logging.getLogger(__name__)
14
14
 
15
15
 
16
- class PaddlePaddleTableClassification(ClassificationModel, OnnxModel):
16
+ class PaddlePaddleTableClassificationModel(ClassificationModel, OnnxModel):
17
+ classes = ["wired", "wireless"]
17
18
  download_options = DownloadOptions(DownloadPlatform.HUGGINGFACE, HF_REPO_ID, "table_cls.onnx")
18
19
 
19
20
  def __call__(self, input: Iterable[NDArray[np.uint8]]) -> list[ClassificationResult]:
@@ -36,8 +37,6 @@ class PaddlePaddleTableClassification(ClassificationModel, OnnxModel):
36
37
 
37
38
  return result
38
39
 
39
- @staticmethod
40
- def postprocess(pred: Sequence[Sequence[float]]) -> list[ClassificationResult]:
41
- return [
42
- ClassificationResult({0: "wired", 1: "wireless"}[np.argmax(p)], max(p)) for p in pred
43
- ]
40
+ @classmethod
41
+ def postprocess(cls, pred: Sequence[Sequence[float]]) -> list[ClassificationResult]:
42
+ return [ClassificationResult(cls.classes[np.argmax(p)], max(p)) for p in pred]
@@ -13,9 +13,12 @@ from ..tasks.base import BaseModel
13
13
  class OnnxModel(BaseModel, ABC):
14
14
  """Base interface for ONNX models."""
15
15
 
16
+ scale = 1 / 255.0
17
+ mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
18
+ std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
19
+
16
20
  def __init__(self, model_path: Optional[Path | str] = None) -> None:
17
- if model_path is None:
18
- model_path = self.download()
21
+ self.model_path = self.download() if model_path is None else Path(model_path)
19
22
 
20
23
  providers_priority = [
21
24
  "CUDAExecutionProvider",
@@ -26,14 +29,10 @@ class OnnxModel(BaseModel, ABC):
26
29
  available_providers = ort.get_available_providers() # type: ignore
27
30
 
28
31
  self.session = ort.InferenceSession(
29
- model_path,
32
+ self.model_path,
30
33
  providers=[p for p in providers_priority if p in available_providers],
31
34
  )
32
35
 
33
- self.scale = 1 / 255.0
34
- self.mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
35
- self.std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
36
-
37
36
  @property
38
37
  def input_shape(self):
39
38
  return self.session.get_inputs()[0].shape[2:] # assuming NCHW
@@ -8,6 +8,7 @@ from ..download import DownloadOptions, download
8
8
  class BaseModel(ABC):
9
9
  """Base interface for models of any type."""
10
10
 
11
+ model_path: Path
11
12
  download_options: Optional[DownloadOptions] = None
12
13
 
13
14
  @abstractmethod
@@ -14,6 +14,8 @@ class ClassificationResult(NamedTuple):
14
14
  class ClassificationModel(BaseModel, ABC):
15
15
  """Base interface for table classification models."""
16
16
 
17
+ classes: list[str]
18
+
17
19
  @abstractmethod
18
20
  def __call__(self, input: Any) -> list[ClassificationResult]:
19
21
  pass
@@ -16,6 +16,8 @@ class DetectionResult(NamedTuple):
16
16
  class DetectionModel(BaseModel, ABC):
17
17
  """Base interface for detection models."""
18
18
 
19
+ classes: list[str] = []
20
+
19
21
  @abstractmethod
20
22
  def __call__(self, input: Any) -> list[Iterator[DetectionResult]]:
21
23
  pass
@@ -0,0 +1,3 @@
1
+ from .paddlepaddle import PaddlePaddleTablePipeline
2
+
3
+ __all__ = [PaddlePaddleTablePipeline]
@@ -0,0 +1,15 @@
1
+ from abc import ABC, abstractmethod
2
+ from pathlib import Path
3
+ from typing import Any, Optional
4
+
5
+
6
+ class BasePipeline(ABC):
7
+ """Base interface for pipelines of any type."""
8
+
9
+ @abstractmethod
10
+ def __init__(self, models_path: Optional[Path | str] = None) -> None:
11
+ pass
12
+
13
+ @abstractmethod
14
+ def __call__(self, input: Any):
15
+ pass
@@ -0,0 +1,62 @@
1
+ import logging
2
+ from abc import ABC, abstractmethod
3
+ from pathlib import Path
4
+ from typing import Any, Iterable, Optional
5
+
6
+ from ..datamodels import Table
7
+ from ..models.tasks import ClassificationModel, DetectionModel
8
+ from .base import BasePipeline
9
+
10
+ logger = logging.getLogger(__name__)
11
+
12
+
13
+ class ClassificationDetectionPipeline(BasePipeline, ABC):
14
+ """Base class for standard classification and detection pipelines."""
15
+
16
+ classification_model: ClassificationModel
17
+ detection_models: list[DetectionModel]
18
+
19
+ @abstractmethod
20
+ def __init__(self, models_path: Optional[Path | str] = None) -> None:
21
+ """Initialize the models."""
22
+
23
+ pass
24
+
25
+ def __call__(self, input: Iterable[Any]) -> list[Table]:
26
+ """Run the pipeline."""
27
+
28
+ cls_images = [[] for c in self.classification_model.classes]
29
+ cls_detections = [[] for c in self.classification_model.classes]
30
+ cls_current_idx = [0 for c in self.classification_model.classes]
31
+ output = []
32
+
33
+ cls_result = self.classification_model(input)
34
+
35
+ # Run the classification model for each image
36
+ for i, (img, p) in enumerate(zip(input, cls_result)):
37
+ cls_images[self.assigned_model_idx(p.cls, self.detection_models)].append(img)
38
+ logger.info("Image %d classified as %s with %.4f confidence", i, p.cls, p.confidence)
39
+
40
+ # Run the detection model for each image
41
+ for i in range(len(self.classification_model.classes)):
42
+ if len(cls_images[i]):
43
+ cls_detections[i] = self.detection_models[i](cls_images[i])
44
+
45
+ # Combine results
46
+ for i in range(len(cls_result)):
47
+ model_idx = self.assigned_model_idx(cls_result[i].cls, self.detection_models)
48
+ cells_det = cls_detections[model_idx][cls_current_idx[model_idx]]
49
+ cls_current_idx[model_idx] += 1
50
+ output.append(Table.from_detections(cells_det))
51
+
52
+ return output
53
+
54
+ @staticmethod
55
+ def assigned_model_idx(pred_cls: str, models: list[DetectionModel]) -> int:
56
+ """Return the index of the first model appropriate for the class."""
57
+
58
+ for idx, model in enumerate(models):
59
+ if pred_cls in model.classes:
60
+ return idx
61
+
62
+ raise ValueError(f"No model can be assigned for class {pred_cls}")
@@ -0,0 +1,45 @@
1
+ import logging
2
+ from pathlib import Path
3
+ from typing import Optional
4
+
5
+ from ..models.PaddlePaddle import (
6
+ PaddlePaddleTableClassificationModel,
7
+ PaddlePaddleWiredCellDetectionModel,
8
+ PaddlePaddleWirelessCellDetectionModel,
9
+ )
10
+ from .classification_detection import ClassificationDetectionPipeline
11
+
12
+ logger = logging.getLogger(__name__)
13
+
14
+
15
+ class PaddlePaddleTablePipeline(ClassificationDetectionPipeline):
16
+ """A table pipeline combining PaddlePaddle classification and detection models."""
17
+
18
+ def __init__(self, models_path: Optional[Path | str] = None) -> None:
19
+ """Initialize models from the provided path or download them.
20
+
21
+ As the models are all in the same repository, do the download only once.
22
+ """
23
+
24
+ self.detection_models = []
25
+ cls_path, wired_path, wireless_path = None, None, None
26
+
27
+ if models_path is not None:
28
+ models_path = Path(models_path)
29
+ cls_path = (
30
+ models_path / PaddlePaddleTableClassificationModel.download_options.model_path
31
+ )
32
+
33
+ self.classification_model = PaddlePaddleTableClassificationModel(cls_path)
34
+
35
+ wired_path = (
36
+ self.classification_model.model_path.parent
37
+ / PaddlePaddleWiredCellDetectionModel.download_options.model_path
38
+ )
39
+ wireless_path = (
40
+ self.classification_model.model_path.parent
41
+ / PaddlePaddleWirelessCellDetectionModel.download_options.model_path
42
+ )
43
+
44
+ self.detection_models.append(PaddlePaddleWiredCellDetectionModel(wired_path))
45
+ self.detection_models.append(PaddlePaddleWirelessCellDetectionModel(wireless_path))
File without changes
File without changes
@@ -2,7 +2,7 @@ import cv2
2
2
  import numpy as np
3
3
  from numpy.typing import NDArray
4
4
 
5
- from .table import Table
5
+ from ..datamodels import Table
6
6
 
7
7
 
8
8
  def visualize_table(
@@ -1,3 +0,0 @@
1
- from .models import DefaultPipeline
2
-
3
- __all__ = [DefaultPipeline]
@@ -1,15 +0,0 @@
1
- from .cell_detection import (
2
- PaddlePaddleCellDetection,
3
- PaddlePaddleWiredCellDetection,
4
- PaddlePaddleWirelessCellDetection,
5
- )
6
- from .pipeline import PaddlePaddleTablePipeline
7
- from .table_classification import PaddlePaddleTableClassification
8
-
9
- __all__ = [
10
- PaddlePaddleCellDetection,
11
- PaddlePaddleWiredCellDetection,
12
- PaddlePaddleWirelessCellDetection,
13
- PaddlePaddleTablePipeline,
14
- PaddlePaddleTableClassification,
15
- ]
@@ -1,59 +0,0 @@
1
- import logging
2
- from pathlib import Path
3
- from typing import Iterable, Optional
4
-
5
- import numpy as np
6
- from numpy.typing import NDArray
7
-
8
- from ...utils.table import Table
9
- from .cell_detection import PaddlePaddleWiredCellDetection, PaddlePaddleWirelessCellDetection
10
- from .table_classification import PaddlePaddleTableClassification
11
-
12
- logger = logging.getLogger(__name__)
13
-
14
-
15
- class PaddlePaddleTablePipeline:
16
- """A table pipeline combining PaddlePaddle classification and detection models."""
17
-
18
- def __init__(self, models_path: Optional[Path | str] = None):
19
- cls_path, wired_path, wireless_path = None, None, None
20
-
21
- if models_path is not None:
22
- models_path = Path(models_path)
23
- cls_path = models_path / PaddlePaddleTableClassification.download_options.model_path
24
- wired_path = models_path / PaddlePaddleWiredCellDetection.download_options.model_path
25
- wireless_path = models_path / PaddlePaddleWiredCellDetection.download_options.model_path
26
-
27
- self.cls_predictor = PaddlePaddleTableClassification(cls_path)
28
- self.wired_predictor = PaddlePaddleWiredCellDetection(wired_path)
29
- self.wireless_predictor = PaddlePaddleWirelessCellDetection(wireless_path)
30
-
31
- def __call__(self, input: Iterable[NDArray[np.uint8]]) -> list[Table]:
32
- wired_images, wireless_images, output = [], [], []
33
-
34
- cls_result = self.cls_predictor(input)
35
-
36
- for i, (img, p) in enumerate(zip(input, cls_result)):
37
- (wired_images if p.cls == "wired" else wireless_images).append(img)
38
- logger.info("Image %d classified as %s", i, p.cls)
39
-
40
- if len(wired_images):
41
- wired_cells = self.wired_predictor(wired_images)
42
-
43
- if len(wireless_images):
44
- wireless_cells = self.wireless_predictor(wireless_images)
45
-
46
- wired_idx = 0
47
- wireless_idx = 0
48
-
49
- for i in range(len(cls_result)):
50
- if cls_result[i].cls == "wired":
51
- cells_det = wired_cells[wired_idx]
52
- wired_idx += 1
53
- else:
54
- cells_det = wireless_cells[wireless_idx]
55
- wireless_idx += 1
56
-
57
- output.append(Table.from_detections(cells_det))
58
-
59
- return output
@@ -1,3 +0,0 @@
1
- from .PaddlePaddle import PaddlePaddleTablePipeline as DefaultPipeline
2
-
3
- __all__ = [DefaultPipeline]
@@ -1,3 +0,0 @@
1
- from .onnx import OnnxModel
2
-
3
- __all__ = [OnnxModel]
File without changes