Update paddle, and support invoice line item
This commit is contained in:
Yaojia Wang
660
tests/table/test_structure_detector.py
Normal file
660
tests/table/test_structure_detector.py
Normal file
@@ -0,0 +1,660 @@
|
||||
"""
|
||||
Tests for PP-StructureV3 Table Detection
|
||||
|
||||
TDD tests for TableDetector class. Tests are designed to run without
|
||||
requiring the actual PP-StructureV3 library by using mock objects.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
import numpy as np
|
||||
|
||||
from backend.table.structure_detector import (
|
||||
TableDetectionResult,
|
||||
TableDetector,
|
||||
TableDetectorConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestTableDetectionResult:
|
||||
"""Tests for TableDetectionResult dataclass."""
|
||||
|
||||
def test_create_with_required_fields(self):
|
||||
"""Test creating result with required fields."""
|
||||
result = TableDetectionResult(
|
||||
bbox=(10.0, 20.0, 300.0, 400.0),
|
||||
html="<table><tr><td>Test</td></tr></table>",
|
||||
confidence=0.95,
|
||||
table_type="wired",
|
||||
)
|
||||
|
||||
assert result.bbox == (10.0, 20.0, 300.0, 400.0)
|
||||
assert result.html == "<table><tr><td>Test</td></tr></table>"
|
||||
assert result.confidence == 0.95
|
||||
assert result.table_type == "wired"
|
||||
assert result.cells == []
|
||||
|
||||
def test_create_with_cells(self):
|
||||
"""Test creating result with cell data."""
|
||||
cells = [
|
||||
{"text": "Header1", "row": 0, "col": 0},
|
||||
{"text": "Value1", "row": 1, "col": 0},
|
||||
]
|
||||
result = TableDetectionResult(
|
||||
bbox=(0, 0, 100, 100),
|
||||
html="<table></table>",
|
||||
confidence=0.9,
|
||||
table_type="wireless",
|
||||
cells=cells,
|
||||
)
|
||||
|
||||
assert len(result.cells) == 2
|
||||
assert result.cells[0]["text"] == "Header1"
|
||||
assert result.table_type == "wireless"
|
||||
|
||||
def test_bbox_is_tuple_of_floats(self):
|
||||
"""Test that bbox contains float values."""
|
||||
result = TableDetectionResult(
|
||||
bbox=(10, 20, 300, 400), # int inputs
|
||||
html="",
|
||||
confidence=0.9,
|
||||
table_type="wired",
|
||||
)
|
||||
|
||||
# Should work with int inputs (duck typing)
|
||||
assert len(result.bbox) == 4
|
||||
|
||||
|
||||
class TestTableDetectorConfig:
|
||||
"""Tests for TableDetectorConfig dataclass."""
|
||||
|
||||
def test_default_values(self):
|
||||
"""Test default configuration values."""
|
||||
config = TableDetectorConfig()
|
||||
|
||||
assert config.device == "gpu:0"
|
||||
assert config.use_doc_orientation_classify is False
|
||||
assert config.use_doc_unwarping is False
|
||||
assert config.use_textline_orientation is False
|
||||
# SLANeXt models for better table recognition accuracy
|
||||
assert config.wired_table_model == "SLANeXt_wired"
|
||||
assert config.wireless_table_model == "SLANeXt_wireless"
|
||||
assert config.layout_model == "PP-DocLayout_plus-L"
|
||||
assert config.min_confidence == 0.5
|
||||
|
||||
def test_custom_values(self):
|
||||
"""Test custom configuration values."""
|
||||
config = TableDetectorConfig(
|
||||
device="cpu",
|
||||
min_confidence=0.7,
|
||||
wired_table_model="SLANet_plus",
|
||||
)
|
||||
|
||||
assert config.device == "cpu"
|
||||
assert config.min_confidence == 0.7
|
||||
assert config.wired_table_model == "SLANet_plus"
|
||||
|
||||
|
||||
class TestTableDetectorInitialization:
|
||||
"""Tests for TableDetector initialization."""
|
||||
|
||||
def test_init_with_default_config(self):
|
||||
"""Test initialization with default config."""
|
||||
detector = TableDetector()
|
||||
|
||||
assert detector.config is not None
|
||||
assert detector.config.device == "gpu:0"
|
||||
assert detector._initialized is False
|
||||
|
||||
def test_init_with_custom_config(self):
|
||||
"""Test initialization with custom config."""
|
||||
config = TableDetectorConfig(device="cpu", min_confidence=0.8)
|
||||
detector = TableDetector(config=config)
|
||||
|
||||
assert detector.config.device == "cpu"
|
||||
assert detector.config.min_confidence == 0.8
|
||||
|
||||
def test_init_with_mock_pipeline(self):
|
||||
"""Test initialization with pre-initialized pipeline."""
|
||||
mock_pipeline = MagicMock()
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
|
||||
assert detector._initialized is True
|
||||
assert detector._pipeline is mock_pipeline
|
||||
|
||||
|
||||
class TestTableDetectorDetection:
|
||||
"""Tests for TableDetector.detect() method."""
|
||||
|
||||
def create_mock_element(
|
||||
self,
|
||||
label: str = "table",
|
||||
bbox: tuple = (10, 20, 300, 400),
|
||||
html: str = "<table><tr><td>Test</td></tr></table>",
|
||||
score: float = 0.95,
|
||||
) -> MagicMock:
|
||||
"""Create a mock PP-StructureV3 element."""
|
||||
element = MagicMock()
|
||||
element.label = label
|
||||
element.bbox = bbox
|
||||
element.html = html
|
||||
element.score = score
|
||||
element.cells = []
|
||||
return element
|
||||
|
||||
def create_mock_result(self, elements: list) -> MagicMock:
|
||||
"""Create a mock PP-StructureV3 result (legacy API without 'get')."""
|
||||
# Use spec=[] to prevent MagicMock from having a 'get' method
|
||||
# This simulates the legacy API that uses layout_elements attribute
|
||||
result = MagicMock(spec=["layout_elements"])
|
||||
result.layout_elements = elements
|
||||
return result
|
||||
|
||||
def test_detect_single_table(self):
|
||||
"""Test detecting a single table in image."""
|
||||
# Setup mock pipeline
|
||||
mock_pipeline = MagicMock()
|
||||
element = self.create_mock_element()
|
||||
mock_result = self.create_mock_result([element])
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].bbox == (10.0, 20.0, 300.0, 400.0)
|
||||
assert results[0].confidence == 0.95
|
||||
assert results[0].table_type == "wired"
|
||||
mock_pipeline.predict.assert_called_once()
|
||||
|
||||
def test_detect_multiple_tables(self):
|
||||
"""Test detecting multiple tables in image."""
|
||||
mock_pipeline = MagicMock()
|
||||
element1 = self.create_mock_element(
|
||||
bbox=(10, 20, 300, 200),
|
||||
html="<table>1</table>",
|
||||
)
|
||||
element2 = self.create_mock_element(
|
||||
bbox=(10, 220, 300, 400),
|
||||
html="<table>2</table>",
|
||||
)
|
||||
mock_result = self.create_mock_result([element1, element2])
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((500, 400, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 2
|
||||
assert results[0].html == "<table>1</table>"
|
||||
assert results[1].html == "<table>2</table>"
|
||||
|
||||
def test_detect_no_tables(self):
|
||||
"""Test handling of image with no tables."""
|
||||
mock_pipeline = MagicMock()
|
||||
# Return result with non-table elements
|
||||
text_element = MagicMock()
|
||||
text_element.label = "text"
|
||||
mock_result = self.create_mock_result([text_element])
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 0
|
||||
|
||||
def test_detect_filters_low_confidence(self):
|
||||
"""Test that low confidence tables are filtered out."""
|
||||
mock_pipeline = MagicMock()
|
||||
low_conf_element = self.create_mock_element(score=0.3)
|
||||
high_conf_element = self.create_mock_element(score=0.9)
|
||||
mock_result = self.create_mock_result([low_conf_element, high_conf_element])
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
config = TableDetectorConfig(min_confidence=0.5)
|
||||
detector = TableDetector(config=config, pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].confidence == 0.9
|
||||
|
||||
def test_detect_wireless_table(self):
|
||||
"""Test detecting wireless (borderless) table."""
|
||||
mock_pipeline = MagicMock()
|
||||
element = self.create_mock_element(label="wireless_table")
|
||||
mock_result = self.create_mock_result([element])
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].table_type == "wireless"
|
||||
|
||||
def test_detect_with_file_path(self):
|
||||
"""Test detection with file path input."""
|
||||
mock_pipeline = MagicMock()
|
||||
element = self.create_mock_element()
|
||||
mock_result = self.create_mock_result([element])
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
|
||||
# Should accept string path
|
||||
results = detector.detect("/path/to/image.png")
|
||||
|
||||
mock_pipeline.predict.assert_called_with("/path/to/image.png")
|
||||
|
||||
def test_detect_returns_empty_on_none_results(self):
|
||||
"""Test handling of None results from pipeline."""
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.predict.return_value = None
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert results == []
|
||||
|
||||
|
||||
class TestTableDetectorLazyInit:
|
||||
"""Tests for lazy initialization of PP-StructureV3."""
|
||||
|
||||
def test_lazy_init_flag_starts_false(self):
|
||||
"""Test that pipeline is not initialized on construction."""
|
||||
detector = TableDetector()
|
||||
assert detector._initialized is False
|
||||
assert detector._pipeline is None
|
||||
|
||||
def test_lazy_init_with_injected_pipeline(self):
|
||||
"""Test that injected pipeline skips lazy initialization."""
|
||||
mock_pipeline = MagicMock()
|
||||
mock_pipeline.predict.return_value = []
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
|
||||
assert detector._initialized is True
|
||||
assert detector._pipeline is mock_pipeline
|
||||
|
||||
# Detection should work without triggering _ensure_initialized import
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
results = detector.detect(image)
|
||||
|
||||
assert results == []
|
||||
mock_pipeline.predict.assert_called_once()
|
||||
|
||||
def test_import_error_without_paddleocr(self):
|
||||
"""Test ImportError when paddleocr is not available."""
|
||||
detector = TableDetector()
|
||||
|
||||
# Simulate paddleocr not being installed
|
||||
with patch.dict("sys.modules", {"paddleocr": None}):
|
||||
with pytest.raises(ImportError) as exc_info:
|
||||
detector._ensure_initialized()
|
||||
|
||||
assert "paddleocr" in str(exc_info.value).lower()
|
||||
|
||||
|
||||
class TestTableDetectorParseResults:
|
||||
"""Tests for result parsing logic."""
|
||||
|
||||
def test_parse_element_with_box_attribute(self):
|
||||
"""Test parsing element with 'box' instead of 'bbox'."""
|
||||
mock_pipeline = MagicMock()
|
||||
element = MagicMock()
|
||||
element.label = "table"
|
||||
element.box = [10, 20, 300, 400] # 'box' instead of 'bbox'
|
||||
element.html = "<table></table>"
|
||||
element.score = 0.9
|
||||
element.cells = []
|
||||
del element.bbox # Remove bbox attribute
|
||||
|
||||
mock_result = MagicMock(spec=["layout_elements"])
|
||||
mock_result.layout_elements = [element]
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].bbox == (10.0, 20.0, 300.0, 400.0)
|
||||
|
||||
def test_parse_element_with_table_html_attribute(self):
|
||||
"""Test parsing element with 'table_html' instead of 'html'."""
|
||||
mock_pipeline = MagicMock()
|
||||
element = MagicMock()
|
||||
element.label = "table"
|
||||
element.bbox = [0, 0, 100, 100]
|
||||
element.table_html = "<table><tr><td>Content</td></tr></table>"
|
||||
element.score = 0.9
|
||||
element.cells = []
|
||||
del element.html
|
||||
|
||||
mock_result = MagicMock(spec=["layout_elements"])
|
||||
mock_result.layout_elements = [element]
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert "<table>" in results[0].html
|
||||
|
||||
def test_parse_element_with_type_attribute(self):
|
||||
"""Test parsing element with 'type' instead of 'label'."""
|
||||
mock_pipeline = MagicMock()
|
||||
element = MagicMock()
|
||||
element.type = "table" # 'type' instead of 'label'
|
||||
element.bbox = [0, 0, 100, 100]
|
||||
element.html = "<table></table>"
|
||||
element.score = 0.9
|
||||
element.cells = []
|
||||
del element.label
|
||||
|
||||
mock_result = MagicMock(spec=["layout_elements"])
|
||||
mock_result.layout_elements = [element]
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
|
||||
def test_parse_cells_data(self):
|
||||
"""Test parsing cell-level data from element."""
|
||||
mock_pipeline = MagicMock()
|
||||
|
||||
# Create mock cells
|
||||
cell1 = MagicMock()
|
||||
cell1.text = "Header"
|
||||
cell1.row = 0
|
||||
cell1.col = 0
|
||||
cell1.row_span = 1
|
||||
cell1.col_span = 1
|
||||
cell1.bbox = [0, 0, 50, 20]
|
||||
|
||||
cell2 = MagicMock()
|
||||
cell2.text = "Value"
|
||||
cell2.row = 1
|
||||
cell2.col = 0
|
||||
cell2.row_span = 1
|
||||
cell2.col_span = 1
|
||||
cell2.bbox = [0, 20, 50, 40]
|
||||
|
||||
element = MagicMock()
|
||||
element.label = "table"
|
||||
element.bbox = [0, 0, 100, 100]
|
||||
element.html = "<table></table>"
|
||||
element.score = 0.9
|
||||
element.cells = [cell1, cell2]
|
||||
|
||||
mock_result = MagicMock(spec=["layout_elements"])
|
||||
mock_result.layout_elements = [element]
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert len(results[0].cells) == 2
|
||||
assert results[0].cells[0]["text"] == "Header"
|
||||
assert results[0].cells[0]["row"] == 0
|
||||
assert results[0].cells[1]["text"] == "Value"
|
||||
assert results[0].cells[1]["row"] == 1
|
||||
|
||||
|
||||
class TestTableDetectorEdgeCases:
|
||||
"""Tests for edge cases and error handling."""
|
||||
|
||||
def test_handles_malformed_element_gracefully(self):
|
||||
"""Test graceful handling of malformed element data."""
|
||||
mock_pipeline = MagicMock()
|
||||
|
||||
# Element missing required attributes
|
||||
bad_element = MagicMock()
|
||||
bad_element.label = "table"
|
||||
# Missing bbox, html, score
|
||||
del bad_element.bbox
|
||||
del bad_element.box
|
||||
|
||||
good_element = MagicMock()
|
||||
good_element.label = "table"
|
||||
good_element.bbox = [0, 0, 100, 100]
|
||||
good_element.html = "<table></table>"
|
||||
good_element.score = 0.9
|
||||
good_element.cells = []
|
||||
|
||||
mock_result = MagicMock(spec=["layout_elements"])
|
||||
mock_result.layout_elements = [bad_element, good_element]
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
# Should not raise, should skip bad element
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
|
||||
def test_handles_empty_layout_elements(self):
|
||||
"""Test handling of empty layout_elements list."""
|
||||
mock_pipeline = MagicMock()
|
||||
mock_result = MagicMock(spec=["layout_elements"])
|
||||
mock_result.layout_elements = []
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert results == []
|
||||
|
||||
def test_handles_result_without_layout_elements(self):
|
||||
"""Test handling of result without layout_elements attribute."""
|
||||
mock_pipeline = MagicMock()
|
||||
mock_result = MagicMock(spec=[]) # No attributes
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert results == []
|
||||
|
||||
def test_confidence_as_list(self):
|
||||
"""Test handling confidence score as list."""
|
||||
mock_pipeline = MagicMock()
|
||||
element = MagicMock()
|
||||
element.label = "table"
|
||||
element.bbox = [0, 0, 100, 100]
|
||||
element.html = "<table></table>"
|
||||
element.score = [0.95] # Score as list
|
||||
element.cells = []
|
||||
|
||||
mock_result = MagicMock(spec=["layout_elements"])
|
||||
mock_result.layout_elements = [element]
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].confidence == 0.95
|
||||
|
||||
|
||||
class TestPaddleX3xAPI:
|
||||
"""Tests for PaddleX 3.x API support (LayoutParsingResultV2)."""
|
||||
|
||||
def test_parse_paddlex_result_with_tables(self):
|
||||
"""Test parsing PaddleX 3.x LayoutParsingResultV2 with tables."""
|
||||
mock_pipeline = MagicMock()
|
||||
|
||||
# Simulate PaddleX 3.x dict-like result
|
||||
mock_result = {
|
||||
"table_res_list": [
|
||||
{
|
||||
"cell_box_list": [[0, 0, 50, 20], [50, 0, 100, 20]],
|
||||
"pred_html": "<table><tr><td>Cell1</td><td>Cell2</td></tr></table>",
|
||||
"table_ocr_pred": ["Cell1", "Cell2"],
|
||||
"table_region_id": 0,
|
||||
}
|
||||
],
|
||||
"parsing_res_list": [
|
||||
{"label": "table", "bbox": [10, 20, 200, 300]},
|
||||
],
|
||||
}
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].html == "<table><tr><td>Cell1</td><td>Cell2</td></tr></table>"
|
||||
assert results[0].bbox == (10.0, 20.0, 200.0, 300.0)
|
||||
assert len(results[0].cells) == 2
|
||||
assert results[0].cells[0]["text"] == "Cell1"
|
||||
assert results[0].cells[1]["text"] == "Cell2"
|
||||
|
||||
def test_parse_paddlex_result_empty_tables(self):
|
||||
"""Test parsing PaddleX 3.x result with no tables."""
|
||||
mock_pipeline = MagicMock()
|
||||
|
||||
mock_result = {
|
||||
"table_res_list": None,
|
||||
"parsing_res_list": [
|
||||
{"label": "text", "bbox": [10, 20, 200, 300]},
|
||||
],
|
||||
}
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 0
|
||||
|
||||
def test_parse_paddlex_result_multiple_tables(self):
|
||||
"""Test parsing PaddleX 3.x result with multiple tables."""
|
||||
mock_pipeline = MagicMock()
|
||||
|
||||
mock_result = {
|
||||
"table_res_list": [
|
||||
{
|
||||
"cell_box_list": [[0, 0, 50, 20]],
|
||||
"pred_html": "<table>1</table>",
|
||||
"table_ocr_pred": ["Text1"],
|
||||
"table_region_id": 0,
|
||||
},
|
||||
{
|
||||
"cell_box_list": [[0, 0, 100, 40]],
|
||||
"pred_html": "<table>2</table>",
|
||||
"table_ocr_pred": ["Text2"],
|
||||
"table_region_id": 1,
|
||||
},
|
||||
],
|
||||
"parsing_res_list": [
|
||||
{"label": "table", "bbox": [10, 20, 200, 300]},
|
||||
{"label": "table", "bbox": [10, 350, 200, 600]},
|
||||
],
|
||||
}
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 2
|
||||
assert results[0].html == "<table>1</table>"
|
||||
assert results[1].html == "<table>2</table>"
|
||||
assert results[0].bbox == (10.0, 20.0, 200.0, 300.0)
|
||||
assert results[1].bbox == (10.0, 350.0, 200.0, 600.0)
|
||||
|
||||
def test_parse_paddlex_result_with_numpy_arrays(self):
|
||||
"""Test parsing PaddleX 3.x result where bbox/cell_box are numpy arrays."""
|
||||
mock_pipeline = MagicMock()
|
||||
|
||||
# Simulate PaddleX 3.x result with numpy arrays (real PP-StructureV3 returns these)
|
||||
mock_result = {
|
||||
"table_res_list": [
|
||||
{
|
||||
"cell_box_list": [
|
||||
np.array([0.0, 0.0, 50.0, 20.0]),
|
||||
np.array([50.0, 0.0, 100.0, 20.0]),
|
||||
],
|
||||
"pred_html": "<table><tr><td>A</td><td>B</td></tr></table>",
|
||||
"table_ocr_pred": ["A", "B"],
|
||||
}
|
||||
],
|
||||
"parsing_res_list": [
|
||||
{"label": "table", "bbox": np.array([10.0, 20.0, 200.0, 300.0])},
|
||||
],
|
||||
}
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].bbox == (10.0, 20.0, 200.0, 300.0)
|
||||
assert results[0].html == "<table><tr><td>A</td><td>B</td></tr></table>"
|
||||
assert len(results[0].cells) == 2
|
||||
assert results[0].cells[0]["text"] == "A"
|
||||
assert results[0].cells[0]["bbox"] == [0.0, 0.0, 50.0, 20.0]
|
||||
assert results[0].cells[1]["text"] == "B"
|
||||
|
||||
def test_parse_paddlex_result_with_empty_numpy_arrays(self):
|
||||
"""Test parsing PaddleX 3.x result where some arrays are empty."""
|
||||
mock_pipeline = MagicMock()
|
||||
|
||||
mock_result = {
|
||||
"table_res_list": [
|
||||
{
|
||||
"cell_box_list": np.array([]), # Empty numpy array
|
||||
"pred_html": "<table></table>",
|
||||
"table_ocr_pred": np.array([]), # Empty numpy array
|
||||
}
|
||||
],
|
||||
"parsing_res_list": [
|
||||
{"label": "table", "bbox": np.array([10.0, 20.0, 200.0, 300.0])},
|
||||
],
|
||||
}
|
||||
mock_pipeline.predict.return_value = [mock_result]
|
||||
|
||||
detector = TableDetector(pipeline=mock_pipeline)
|
||||
image = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
|
||||
results = detector.detect(image)
|
||||
|
||||
assert len(results) == 1
|
||||
assert results[0].cells == [] # Empty cells list
|
||||
assert results[0].html == "<table></table>"
|
||||
Reference in New Issue
Block a user