agent_jrxml/tests/test_annotation_detector.py

"""测试批注检测器：圆圈检测、箭头检测、OCR 关联、格式化。"""

import tempfile
from pathlib import Path

import cv2
import numpy as np
import pytest


def _draw_circle_image(path: str, size: tuple = (400, 300)) -> None:
    """生成包含红色圆圈的合成测试图片。"""
    img = np.ones((size[1], size[0], 3), dtype=np.uint8) * 255
    cv2.circle(img, (200, 150), 50, (0, 0, 255), 2)
    cv2.imwrite(path, img)


def _draw_arrow_image(path: str, size: tuple = (400, 300)) -> None:
    """生成包含手绘风格箭头的合成测试图片（多段线模拟手绘）。"""
    img = np.ones((size[1], size[0], 3), dtype=np.uint8) * 255
    # 多段略微偏移的线段模拟手绘箭杆（产生多个 HoughLinesP 段）
    for offset_y in (-1, 0, 1):
        cv2.line(img, (50, 150 + offset_y), (200, 150 + offset_y), (0, 0, 255), 2)
    for offset_y in (-1, 0, 1):
        cv2.line(img, (200, 150 + offset_y), (340, 150 + offset_y), (0, 0, 255), 2)
    # 箭头三角形
    pts = np.array([[350, 150], [330, 135], [330, 165]], np.int32)
    cv2.fillPoly(img, [pts], (0, 0, 255))
    # 额外三角形边缘线
    cv2.line(img, (350, 150), (330, 135), (0, 0, 255), 2)
    cv2.line(img, (350, 150), (330, 165), (0, 0, 255), 2)
    cv2.imwrite(path, img)


def _draw_circle_and_text_image(path: str, size: tuple = (500, 400)) -> None:
    """生成包含红色圆圈和"文本"的合成图片（模拟圈选批注）。"""
    img = np.ones((size[1], size[0], 3), dtype=np.uint8) * 255
    cv2.circle(img, (250, 150), 60, (0, 0, 255), 3)
    cv2.putText(img, "项目A", (20, 160), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 0), 2)
    cv2.imwrite(path, img)


class TestAnnotationDetector:
    """测试 annotation_detector.py 各功能。"""

    def test_detect_circles_finds_circle(self):
        from backend.annotation_detector import detect_annotations

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            path = tmp.name
        try:
            _draw_circle_image(path)
            ocr_elements = [
                {"text": "测试字段", "bbox": {"x": 170, "y": 120, "w": 60, "h": 20}, "confidence": 0.95},
            ]
            result = detect_annotations(path, ocr_elements)
            assert result["total"] >= 1
            circles = result["circles"]
            assert len(circles) >= 1
            c = circles[0]
            assert c["type"] == "circle"
            assert "center" in c
            assert "bbox" in c
            assert "nearby_texts" in c
        finally:
            Path(path).unlink(missing_ok=True)

    def test_detect_arrows_finds_arrow(self):
        from backend.annotation_detector import detect_annotations

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            path = tmp.name
        try:
            _draw_arrow_image(path)
            ocr_elements = [
                {"text": "起点", "bbox": {"x": 30, "y": 130, "w": 40, "h": 20}, "confidence": 0.9},
                {"text": "终点", "bbox": {"x": 310, "y": 130, "w": 40, "h": 20}, "confidence": 0.9},
            ]
            result = detect_annotations(path, ocr_elements)
            assert result["total"] >= 1
            arrows = result["arrows"]
            assert len(arrows) >= 1
            a = arrows[0]
            assert a["type"] == "arrow"
            assert "from_pt" in a
            assert "to_pt" in a
        finally:
            Path(path).unlink(missing_ok=True)

    def test_correlate_with_ocr_links_nearby_texts(self):
        from backend.annotation_detector import detect_annotations

        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            path = tmp.name
        try:
            _draw_circle_and_text_image(path)
            ocr_elements = [
                {"text": "项目A", "bbox": {"x": 20, "y": 140, "w": 80, "h": 30}, "confidence": 0.98},
                {"text": "金额", "bbox": {"x": 350, "y": 200, "w": 50, "h": 20}, "confidence": 0.9},
            ]
            result = detect_annotations(path, ocr_elements)
            circles = result["circles"]
            if circles:
                near = circles[0].get("nearby_texts", [])
                if near:
                    assert "项目A" in near
        finally:
            Path(path).unlink(missing_ok=True)

    def test_invalid_image_path(self):
        from backend.annotation_detector import detect_annotations

        result = detect_annotations("/nonexistent/file.png", [])
        assert result["total"] == 0
        assert "error" in result

    def test_format_annotation_context_empty(self):
        from backend.annotation_detector import format_annotation_context

        assert format_annotation_context({}) == ""
        assert format_annotation_context(None) == ""
        assert format_annotation_context({"circles": [], "arrows": [], "total": 0}) == ""

    def test_format_annotation_context_with_circles(self):
        from backend.annotation_detector import format_annotation_context

        ann = {
            "circles": [
                {"center": [100, 200], "nearby_texts": ["项目A", "金额"]},
            ],
            "arrows": [],
            "total": 1,
        }
        text = format_annotation_context(ann)
        assert "圈选标记" in text
        assert "项目A" in text

    def test_format_annotation_context_with_arrows(self):
        from backend.annotation_detector import format_annotation_context

        ann = {
            "circles": [],
            "arrows": [
                {"from_text": "修理号", "to_text": "车架号"},
            ],
            "total": 1,
        }
        text = format_annotation_context(ann)
        assert "箭头标记" in text
        assert "修理号" in text
        assert "车架号" in text