84 lines
3.0 KiB
Python
84 lines
3.0 KiB
Python
from __future__ import annotations
|
|
|
|
from PIL import Image, ImageDraw
|
|
|
|
from api.services.docling import extract
|
|
from api.services.docling.regions import detect_response_regions_from_image
|
|
|
|
|
|
def test_detects_grouped_answer_lines() -> None:
|
|
image = Image.new("RGB", (900, 1200), "white")
|
|
draw = ImageDraw.Draw(image)
|
|
for y in (420, 470, 520):
|
|
draw.line((160, y, 760, y), fill="black", width=3)
|
|
|
|
candidates = detect_response_regions_from_image(image, page_index=2)
|
|
|
|
line_regions = [c.to_mapper_dict() for c in candidates if c.region_type == "answer_lines"]
|
|
assert line_regions
|
|
best = line_regions[0]
|
|
assert best["kind"] == "response"
|
|
assert best["source"] == "ai"
|
|
assert best["confirmed"] is False
|
|
assert best["page_index"] == 2
|
|
assert best["line_count"] == 3
|
|
assert best["bbox"]["coord_origin"] == "TOPLEFT"
|
|
assert best["bbox"]["w"] > 550
|
|
assert best["bbox"]["h"] > 80
|
|
|
|
|
|
def test_detects_answer_box() -> None:
|
|
image = Image.new("RGB", (900, 1200), "white")
|
|
draw = ImageDraw.Draw(image)
|
|
draw.rectangle((140, 300, 780, 520), outline="black", width=3)
|
|
|
|
candidates = detect_response_regions_from_image(image, page_index=0)
|
|
|
|
boxes = [c.to_mapper_dict() for c in candidates if c.region_type == "answer_box"]
|
|
assert boxes
|
|
assert boxes[0]["bbox"]["w"] > 600
|
|
assert boxes[0]["bbox"]["h"] > 200
|
|
|
|
|
|
def test_detect_response_region_taxonomy_for_lines_and_boxes():
|
|
img = Image.new("RGB", (800, 1000), "white")
|
|
draw = ImageDraw.Draw(img)
|
|
for y in (220, 260, 300):
|
|
draw.line((120, y, 680, y), fill="black", width=2)
|
|
draw.rectangle((140, 520, 660, 640), outline="black", width=3)
|
|
|
|
regions = detect_response_regions_from_image(img, min_confidence=0.1)
|
|
types = {r.region_type for r in regions}
|
|
|
|
assert "answer_lines" in types
|
|
assert "answer_box" in types
|
|
|
|
|
|
def test_attach_detected_response_regions_normalizes_taxonomy_and_uses_pdf_y(monkeypatch, tmp_path):
|
|
pdf = tmp_path / "paper.pdf"
|
|
pdf.write_bytes(b"%PDF test placeholder")
|
|
parts = {
|
|
"01.1": {"q": "01", "page": 1, "bbox": {"l": 50, "r": 80, "t": 700, "b": 680}, "regions": []},
|
|
"01.2": {"q": "01", "page": 1, "bbox": {"l": 50, "r": 80, "t": 500, "b": 480}, "regions": []},
|
|
}
|
|
|
|
def fake_detect(path, min_confidence=0.32):
|
|
return [{
|
|
"page_index": 0,
|
|
"region_type": "answer-box",
|
|
"confidence": 0.77,
|
|
"bbox": {"x": 100, "y": 335, "w": 500, "h": 40},
|
|
"detection_method": "test",
|
|
"meta": {"page_height_px": 1000, "page_height_pdf": 800},
|
|
}]
|
|
|
|
monkeypatch.setattr(extract.region_mod, "detect_response_regions_from_pdf", fake_detect)
|
|
|
|
attached, candidates = extract.attach_detected_response_regions(parts, str(pdf))
|
|
|
|
assert attached == 1
|
|
assert len(candidates) == 1
|
|
assert parts["01.1"]["regions"] == []
|
|
assert parts["01.2"]["regions"][0]["type"] == "answer_box"
|
|
assert parts["01.2"]["regions"][0]["source"] == "opencv"
|