q792602257
41006b098a
Tests / Test apps.device-host-agent.tests.test_e2e.test_public_sdk_reports_fake_device_success_and_runtime_failure failed
PaddleOCR itself returns no color info, only text/bounds/confidence.
Add pixel-level post-processing in perception/ocr.py: crop the
screenshot to each OCR box, split pixels into two luminance clusters
via Otsu threshold, and treat the minority cluster as the text stroke
(foreground) and the majority as the background. New
SceneElement.foreground_color/background_color fields ("#rrggbb",
None when not OCR-sourced or sampling fails) round-trip through
to_dict/from_dict alongside the existing accessibility-state fields.
Planner system prompt documents the new fields as a secondary signal.
pillow is promoted from an implicit paddleocr transitive dependency to
an explicit direct dependency since perception/ocr.py now imports PIL
directly; uv.lock re-resolved with no version change (already locked
at 12.3.0).
148 lines
4.5 KiB
Python
148 lines
4.5 KiB
Python
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from core.models import Bounds, SceneElement
|
|
from perception.ocr import (
|
|
OCRBox,
|
|
PaddleOCREngine,
|
|
_with_sampled_colors,
|
|
parse_paddle_result,
|
|
run_ocr,
|
|
)
|
|
|
|
|
|
def test_paddle_ocr_engine_disables_doc_preprocessing_by_default() -> None:
|
|
# Screenshots are flat, upright digital captures, not photographed paper
|
|
# documents: PaddleOCR's document-unwarping preprocessing has nothing real
|
|
# to correct and instead warps the image, so detected box coordinates no
|
|
# longer line up with the original screenshot.
|
|
engine = PaddleOCREngine(lang="ch")
|
|
|
|
assert engine.kwargs["use_doc_orientation_classify"] is False
|
|
assert engine.kwargs["use_doc_unwarping"] is False
|
|
|
|
|
|
def test_paddle_ocr_engine_lets_callers_override_doc_preprocessing() -> None:
|
|
engine = PaddleOCREngine(use_doc_unwarping=True)
|
|
|
|
assert engine.kwargs["use_doc_unwarping"] is True
|
|
|
|
|
|
def test_parse_paddle_result_accepts_ndarray_fields() -> None:
|
|
boxes = parse_paddle_result(
|
|
{
|
|
"rec_texts": np.array(["Search"]),
|
|
"rec_scores": np.array([0.95]),
|
|
"rec_boxes": np.array([[1, 2, 10, 12]]),
|
|
}
|
|
)
|
|
|
|
assert boxes == [
|
|
OCRBox(
|
|
text="Search",
|
|
bounds=Bounds(x=1.0, y=2.0, width=9.0, height=10.0),
|
|
confidence=0.95,
|
|
)
|
|
]
|
|
|
|
|
|
class CallableJsonResult:
|
|
def json(self) -> dict[str, object]:
|
|
return {
|
|
"res": {
|
|
"rec_texts": ["Search"],
|
|
"rec_scores": [0.95],
|
|
"rec_boxes": [[1, 2, 10, 12]],
|
|
}
|
|
}
|
|
|
|
|
|
def test_parse_paddle_result_accepts_callable_json_result() -> None:
|
|
assert parse_paddle_result(CallableJsonResult()) == [
|
|
OCRBox(
|
|
text="Search",
|
|
bounds=Bounds(x=1.0, y=2.0, width=9.0, height=10.0),
|
|
confidence=0.95,
|
|
)
|
|
]
|
|
|
|
|
|
class FailingOCREngine:
|
|
def extract(self, image: bytes | str) -> list[OCRBox]:
|
|
raise ValueError(
|
|
"The truth value of an array with more than one element is ambiguous"
|
|
)
|
|
|
|
|
|
def test_run_ocr_degrades_when_ocr_engine_raises() -> None:
|
|
assert run_ocr(b"image", engine=FailingOCREngine()) == [] # type: ignore[arg-type]
|
|
|
|
|
|
def test_run_ocr_strict_mode_preserves_engine_error() -> None:
|
|
with pytest.raises(ValueError, match="truth value"):
|
|
run_ocr(b"image", engine=FailingOCREngine(), strict=True) # type: ignore[arg-type]
|
|
|
|
|
|
def test_ocr_box_to_scene_element_carries_sampled_colors() -> None:
|
|
element = OCRBox(
|
|
text="Search",
|
|
bounds=Bounds(x=0, y=0, width=10, height=10),
|
|
foreground_color="#000000",
|
|
background_color="#ffffff",
|
|
).to_scene_element("ocr-000")
|
|
|
|
assert element.foreground_color == "#000000"
|
|
assert element.background_color == "#ffffff"
|
|
data = element.to_dict()
|
|
assert data["foreground_color"] == "#000000"
|
|
assert data["background_color"] == "#ffffff"
|
|
assert SceneElement.from_dict(data).foreground_color == "#000000"
|
|
|
|
|
|
def test_ocr_box_to_scene_element_omits_colors_when_unknown() -> None:
|
|
element = OCRBox(
|
|
text="Search", bounds=Bounds(x=0, y=0, width=10, height=10)
|
|
).to_scene_element("ocr-000")
|
|
|
|
data = element.to_dict()
|
|
|
|
assert "foreground_color" not in data
|
|
assert "background_color" not in data
|
|
|
|
|
|
def test_with_sampled_colors_estimates_foreground_and_background(tmp_path) -> None:
|
|
image_path = tmp_path / "shot.png"
|
|
picture = Image.new("RGB", (20, 20), (255, 255, 255))
|
|
for x in range(8, 12):
|
|
for y in range(8, 12):
|
|
picture.putpixel((x, y), (0, 0, 0))
|
|
picture.save(image_path)
|
|
boxes = [OCRBox(text="A", bounds=Bounds(x=0, y=0, width=20, height=20))]
|
|
|
|
sampled = _with_sampled_colors(boxes, str(image_path))
|
|
|
|
assert sampled[0].foreground_color == "#000000"
|
|
assert sampled[0].background_color == "#ffffff"
|
|
|
|
|
|
def test_with_sampled_colors_skips_degenerate_bounds(tmp_path) -> None:
|
|
image_path = tmp_path / "shot.png"
|
|
Image.new("RGB", (20, 20), (255, 255, 255)).save(image_path)
|
|
boxes = [OCRBox(text="A", bounds=Bounds(x=0, y=0, width=1, height=1))]
|
|
|
|
sampled = _with_sampled_colors(boxes, str(image_path))
|
|
|
|
assert sampled[0].foreground_color is None
|
|
assert sampled[0].background_color is None
|
|
|
|
|
|
def test_with_sampled_colors_degrades_when_image_cannot_be_opened(tmp_path) -> None:
|
|
boxes = [OCRBox(text="A", bounds=Bounds(x=0, y=0, width=20, height=20))]
|
|
|
|
sampled = _with_sampled_colors(boxes, str(tmp_path / "missing.png"))
|
|
|
|
assert sampled == boxes
|