40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
from __future__ import annotations
|
|
|
|
from PIL import Image, ImageDraw
|
|
|
|
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
|