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.
- {cells2table-0.2.0 → cells2table-0.2.2}/PKG-INFO +23 -21
- {cells2table-0.2.0 → cells2table-0.2.2}/README.md +10 -9
- {cells2table-0.2.0 → cells2table-0.2.2}/pyproject.toml +30 -15
- cells2table-0.2.2/src/cells2table/__init__.py +3 -0
- cells2table-0.2.2/src/cells2table/datamodels/__init__.py +4 -0
- cells2table-0.2.2/src/cells2table/datamodels/bbox.py +21 -0
- {cells2table-0.2.0/src/cells2table/utils → cells2table-0.2.2/src/cells2table/datamodels}/table.py +2 -19
- {cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/docling.py +18 -19
- cells2table-0.2.2/src/cells2table/models/PaddlePaddle/__init__.py +13 -0
- {cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/models/PaddlePaddle/cell_detection.py +10 -7
- {cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/models/PaddlePaddle/table_classification.py +8 -9
- {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/runtimes/onnx.py +6 -7
- {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/tasks/base.py +1 -0
- {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/tasks/classification.py +2 -0
- {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/tasks/detection.py +2 -0
- cells2table-0.2.2/src/cells2table/pipelines/__init__.py +3 -0
- cells2table-0.2.2/src/cells2table/pipelines/base.py +15 -0
- cells2table-0.2.2/src/cells2table/pipelines/classification_detection.py +62 -0
- cells2table-0.2.2/src/cells2table/pipelines/paddlepaddle.py +45 -0
- cells2table-0.2.2/src/cells2table/py.typed +0 -0
- cells2table-0.2.2/src/cells2table/utils/__init__.py +0 -0
- {cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/utils/visualize.py +1 -1
- cells2table-0.2.0/src/cells2table/__init__.py +0 -3
- cells2table-0.2.0/src/cells2table/models/PaddlePaddle/__init__.py +0 -15
- cells2table-0.2.0/src/cells2table/models/PaddlePaddle/pipeline.py +0 -59
- cells2table-0.2.0/src/cells2table/models/__init__.py +0 -3
- cells2table-0.2.0/src/cells2table/models/utils/runtimes/__init__.py +0 -3
- {cells2table-0.2.0 → cells2table-0.2.2}/LICENSE +0 -0
- {cells2table-0.2.0/src/cells2table/utils → cells2table-0.2.2/src/cells2table/models}/__init__.py +0 -0
- {cells2table-0.2.0/src/cells2table/models/utils → cells2table-0.2.2/src/cells2table/models}/download.py +0 -0
- /cells2table-0.2.0/src/cells2table/py.typed → /cells2table-0.2.2/src/cells2table/models/runtimes/__init__.py +0 -0
- {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.
|
|
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:
|
|
13
|
-
Requires-Dist:
|
|
14
|
-
Requires-Dist:
|
|
15
|
-
Requires-Dist:
|
|
16
|
-
Requires-Dist: onnxruntime
|
|
17
|
-
Requires-Dist:
|
|
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:
|
|
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
|
|
35
|
+
### PaddlePaddle
|
|
35
36
|
|
|
36
37
|
- Classification model (wired / wireless)
|
|
37
38
|
- Cell detection model with different weights for each class
|
|
38
39
|
|
|
39
|
-
|
|
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
|
|
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
|
|
52
|
-
|
|
|
53
|
-
| `
|
|
54
|
-
| `
|
|
55
|
-
| `
|
|
56
|
-
| `
|
|
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
|
|
7
|
+
### PaddlePaddle
|
|
8
8
|
|
|
9
9
|
- Classification model (wired / wireless)
|
|
10
10
|
- Cell detection model with different weights for each class
|
|
11
11
|
|
|
12
|
-
|
|
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
|
|
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
|
|
25
|
-
|
|
|
26
|
-
| `
|
|
27
|
-
| `
|
|
28
|
-
| `
|
|
29
|
-
| `
|
|
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.
|
|
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
|
-
"
|
|
19
|
-
"
|
|
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 = "
|
|
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
|
-
|
|
35
|
-
"
|
|
33
|
+
docling = [
|
|
34
|
+
"docling (>=2.66.0,<3.0.0)",
|
|
36
35
|
]
|
|
37
|
-
|
|
38
|
-
"
|
|
36
|
+
huggingface = [
|
|
37
|
+
"huggingface-hub (>=0.36.0,<1.0.0)",
|
|
39
38
|
]
|
|
40
|
-
|
|
41
|
-
"
|
|
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
|
-
|
|
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 = [
|
|
65
|
+
conflicts = [
|
|
66
|
+
[
|
|
67
|
+
{ extra = "onnx_cpu" },
|
|
68
|
+
{ extra = "onnx_cuda" },
|
|
69
|
+
{ extra = "onnx_openvino" }
|
|
70
|
+
]
|
|
71
|
+
]
|
|
@@ -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)
|
{cells2table-0.2.0/src/cells2table/utils → cells2table-0.2.2/src/cells2table/datamodels}/table.py
RENAMED
|
@@ -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
|
-
|
|
7
|
-
from
|
|
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
|
-
|
|
22
|
-
from .
|
|
23
|
-
|
|
24
|
-
|
|
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
|
+
]
|
{cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/models/PaddlePaddle/cell_detection.py
RENAMED
|
@@ -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 ..
|
|
8
|
-
from ..
|
|
9
|
-
from ..
|
|
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
|
|
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
|
-
|
|
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
|
|
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
|
|
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
|
)
|
{cells2table-0.2.0 → cells2table-0.2.2}/src/cells2table/models/PaddlePaddle/table_classification.py
RENAMED
|
@@ -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 ..
|
|
8
|
-
from ..
|
|
9
|
-
from ..
|
|
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
|
|
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
|
-
@
|
|
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
|
|
@@ -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
|
|
@@ -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
|
|
@@ -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
|
|
File without changes
|
{cells2table-0.2.0/src/cells2table/utils → cells2table-0.2.2/src/cells2table/models}/__init__.py
RENAMED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|