47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
|
import numpy as np
|
||
|
|
|
||
|
|
from app.services.layout_detector import LayoutDetector
|
||
|
|
|
||
|
|
|
||
|
|
class _FakePredictor:
|
||
|
|
def __init__(self, boxes):
|
||
|
|
self._boxes = boxes
|
||
|
|
|
||
|
|
def predict(self, image):
|
||
|
|
return [{"boxes": self._boxes}]
|
||
|
|
|
||
|
|
|
||
|
|
def test_detect_applies_postprocess_and_keeps_native_label(monkeypatch):
|
||
|
|
raw_boxes = [
|
||
|
|
{"cls_id": 22, "label": "text", "score": 0.95, "coordinate": [0, 0, 100, 100]},
|
||
|
|
{"cls_id": 22, "label": "text", "score": 0.90, "coordinate": [10, 10, 20, 20]},
|
||
|
|
{"cls_id": 6, "label": "doc_title", "score": 0.99, "coordinate": [0, 0, 80, 20]},
|
||
|
|
]
|
||
|
|
|
||
|
|
detector = LayoutDetector.__new__(LayoutDetector)
|
||
|
|
detector._get_layout_detector = lambda: _FakePredictor(raw_boxes)
|
||
|
|
|
||
|
|
calls = {}
|
||
|
|
|
||
|
|
def fake_apply_layout_postprocess(boxes, img_size, layout_nms, layout_unclip_ratio, layout_merge_bboxes_mode):
|
||
|
|
calls["args"] = {
|
||
|
|
"boxes": boxes,
|
||
|
|
"img_size": img_size,
|
||
|
|
"layout_nms": layout_nms,
|
||
|
|
"layout_unclip_ratio": layout_unclip_ratio,
|
||
|
|
"layout_merge_bboxes_mode": layout_merge_bboxes_mode,
|
||
|
|
}
|
||
|
|
return [boxes[0], boxes[2]]
|
||
|
|
|
||
|
|
monkeypatch.setattr("app.services.layout_detector.apply_layout_postprocess", fake_apply_layout_postprocess)
|
||
|
|
|
||
|
|
image = np.zeros((200, 100, 3), dtype=np.uint8)
|
||
|
|
info = detector.detect(image)
|
||
|
|
|
||
|
|
assert calls["args"]["img_size"] == (100, 200)
|
||
|
|
assert calls["args"]["layout_nms"] is True
|
||
|
|
assert calls["args"]["layout_merge_bboxes_mode"] == "large"
|
||
|
|
assert [region.native_label for region in info.regions] == ["text", "doc_title"]
|
||
|
|
assert [region.type for region in info.regions] == ["text", "text"]
|
||
|
|
assert info.MixedRecognition is True
|