doc_processer/tests/services/test_layout_detector.py

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