doc_processer/tests/services/test_layout_postprocess.py

import math

import numpy as np

from app.services.layout_postprocess import (
    apply_layout_postprocess,
    check_containment,
    iou,
    is_contained,
    nms,
    unclip_boxes,
)


def _raw_box(cls_id, score, x1, y1, x2, y2, label="text"):
    return {
        "cls_id": cls_id,
        "label": label,
        "score": score,
        "coordinate": [x1, y1, x2, y2],
    }


def test_iou_handles_full_none_and_partial_overlap():
    assert iou([0, 0, 9, 9], [0, 0, 9, 9]) == 1.0
    assert iou([0, 0, 9, 9], [20, 20, 29, 29]) == 0.0
    assert math.isclose(iou([0, 0, 9, 9], [5, 5, 14, 14]), 1 / 7, rel_tol=1e-6)


def test_nms_keeps_highest_score_for_same_class_overlap():
    boxes = np.array(
        [
            [0, 0.95, 0, 0, 10, 10],
            [0, 0.80, 1, 1, 11, 11],
        ],
        dtype=float,
    )

    kept = nms(boxes, iou_same=0.6, iou_diff=0.98)

    assert kept == [0]


def test_nms_keeps_cross_class_overlap_boxes_below_diff_threshold():
    boxes = np.array(
        [
            [0, 0.95, 0, 0, 10, 10],
            [1, 0.90, 1, 1, 11, 11],
        ],
        dtype=float,
    )

    kept = nms(boxes, iou_same=0.6, iou_diff=0.98)

    assert kept == [0, 1]


def test_nms_returns_single_box_index():
    boxes = np.array([[0, 0.95, 0, 0, 10, 10]], dtype=float)

    assert nms(boxes) == [0]


def test_is_contained_uses_overlap_threshold():
    outer = [0, 0.9, 0, 0, 10, 10]
    inner = [0, 0.9, 2, 2, 8, 8]
    partial = [0, 0.9, 6, 6, 12, 12]

    assert is_contained(inner, outer) is True
    assert is_contained(partial, outer) is False
    assert is_contained(partial, outer, overlap_threshold=0.3) is True


def test_check_containment_respects_preserve_class_ids():
    boxes = np.array(
        [
            [0, 0.9, 0, 0, 100, 100],
            [1, 0.8, 10, 10, 30, 30],
            [2, 0.7, 15, 15, 25, 25],
        ],
        dtype=float,
    )

    contains_other, contained_by_other = check_containment(boxes, preserve_cls_ids={1})

    assert contains_other.tolist() == [1, 1, 0]
    assert contained_by_other.tolist() == [0, 0, 1]


def test_unclip_boxes_supports_scalar_tuple_dict_and_none():
    boxes = np.array(
        [
            [0, 0.9, 10, 10, 20, 20],
            [1, 0.8, 30, 30, 50, 40],
        ],
        dtype=float,
    )

    scalar = unclip_boxes(boxes, 2.0)
    assert scalar[:, 2:6].tolist() == [[5.0, 5.0, 25.0, 25.0], [20.0, 25.0, 60.0, 45.0]]

    tuple_ratio = unclip_boxes(boxes, (2.0, 3.0))
    assert tuple_ratio[:, 2:6].tolist() == [[5.0, 0.0, 25.0, 30.0], [20.0, 20.0, 60.0, 50.0]]

    per_class = unclip_boxes(boxes, {1: (1.5, 2.0)})
    assert per_class[:, 2:6].tolist() == [[10.0, 10.0, 20.0, 20.0], [25.0, 25.0, 55.0, 45.0]]

    assert np.array_equal(unclip_boxes(boxes, None), boxes)


def test_apply_layout_postprocess_large_mode_removes_contained_small_box():
    boxes = [
        _raw_box(0, 0.95, 0, 0, 100, 100, "text"),
        _raw_box(0, 0.90, 10, 10, 20, 20, "text"),
    ]

    result = apply_layout_postprocess(boxes, img_size=(120, 120), layout_merge_bboxes_mode="large")

    assert [box["coordinate"] for box in result] == [[0, 0, 100, 100]]


def test_apply_layout_postprocess_preserves_contained_image_like_boxes():
    boxes = [
        _raw_box(0, 0.95, 0, 0, 100, 100, "text"),
        _raw_box(1, 0.90, 10, 10, 20, 20, "image"),
        _raw_box(2, 0.90, 25, 25, 35, 35, "seal"),
        _raw_box(3, 0.90, 40, 40, 50, 50, "chart"),
    ]

    result = apply_layout_postprocess(boxes, img_size=(120, 120), layout_merge_bboxes_mode="large")

    assert {box["label"] for box in result} == {"text", "image", "seal", "chart"}


def test_apply_layout_postprocess_clamps_skips_invalid_and_filters_large_image():
    boxes = [
        _raw_box(0, 0.95, -10, -5, 40, 50, "text"),
        _raw_box(1, 0.90, 10, 10, 10, 50, "text"),
        _raw_box(2, 0.85, 0, 0, 100, 90, "image"),
    ]

    result = apply_layout_postprocess(
        boxes,
        img_size=(100, 90),
        layout_nms=False,
        layout_merge_bboxes_mode=None,
    )

    assert result == [{"cls_id": 0, "label": "text", "score": 0.95, "coordinate": [0, 0, 40, 50]}]