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).
328 lines
10 KiB
Python
328 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import tempfile
|
|
import logging
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass, replace
|
|
from numbers import Real
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from core.models import Bounds, SceneElement
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OCRBox:
|
|
text: str
|
|
bounds: Bounds
|
|
confidence: float | None = None
|
|
foreground_color: str | None = None
|
|
background_color: str | None = None
|
|
|
|
def to_scene_element(self, element_id: str) -> SceneElement:
|
|
return SceneElement(
|
|
id=element_id,
|
|
type="text",
|
|
text=self.text,
|
|
bounds=self.bounds,
|
|
confidence=self.confidence,
|
|
source="ocr",
|
|
foreground_color=self.foreground_color,
|
|
background_color=self.background_color,
|
|
)
|
|
|
|
|
|
class PaddleOCREngine:
|
|
def __init__(self, **kwargs: Any) -> None:
|
|
# Device screenshots are already upright, flat digital captures, not
|
|
# photographed paper documents, so PaddleOCR's document-preprocessing
|
|
# models (orientation classification + UVDoc unwarping) have nothing
|
|
# real to correct. Left at their library defaults (True), they still
|
|
# run, geometrically warp the image, and detect/recognize text against
|
|
# that warped image, returning box coordinates that no longer line up
|
|
# with the original screenshot. Callers can still opt back in via an
|
|
# explicit kwarg.
|
|
kwargs.setdefault("use_doc_orientation_classify", False)
|
|
kwargs.setdefault("use_doc_unwarping", False)
|
|
self.kwargs = kwargs
|
|
self._engine: Any | None = None
|
|
|
|
def extract(self, image: bytes | str | Path) -> list[OCRBox]:
|
|
engine = self._load()
|
|
image_input, temp_path = _image_input(image)
|
|
try:
|
|
if hasattr(engine, "predict"):
|
|
raw = engine.predict(input=image_input)
|
|
else:
|
|
raw = engine.ocr(image_input, cls=True)
|
|
boxes = parse_paddle_result(raw)
|
|
return _with_sampled_colors(boxes, image_input)
|
|
finally:
|
|
if temp_path:
|
|
temp_path.unlink(missing_ok=True)
|
|
|
|
def _load(self) -> Any:
|
|
if self._engine is None:
|
|
from paddleocr import PaddleOCR
|
|
|
|
self._engine = PaddleOCR(**self.kwargs)
|
|
return self._engine
|
|
|
|
|
|
def run_ocr(
|
|
image: bytes | str | Path,
|
|
*,
|
|
engine: PaddleOCREngine | None = None,
|
|
strict: bool = False,
|
|
) -> list[SceneElement]:
|
|
try:
|
|
boxes = (engine or PaddleOCREngine()).extract(image)
|
|
except Exception:
|
|
if strict:
|
|
raise
|
|
logger.warning("OCR extraction failed; continuing without OCR", exc_info=True)
|
|
return []
|
|
return [box.to_scene_element(f"ocr-{index:03d}") for index, box in enumerate(boxes)]
|
|
|
|
|
|
def _with_sampled_colors(boxes: list[OCRBox], image_path: str) -> list[OCRBox]:
|
|
if not boxes:
|
|
return boxes
|
|
try:
|
|
from PIL import Image
|
|
|
|
with Image.open(image_path) as source:
|
|
picture = source.convert("RGB")
|
|
sampled = []
|
|
for box in boxes:
|
|
foreground, background = _estimate_colors(picture, box.bounds)
|
|
sampled.append(
|
|
replace(
|
|
box,
|
|
foreground_color=foreground,
|
|
background_color=background,
|
|
)
|
|
)
|
|
return sampled
|
|
except Exception:
|
|
logger.warning(
|
|
"OCR color sampling failed; continuing without foreground/background colors",
|
|
exc_info=True,
|
|
)
|
|
return boxes
|
|
|
|
|
|
def _estimate_colors(picture: Any, bounds: Bounds) -> tuple[str | None, str | None]:
|
|
left = max(0, int(bounds.x))
|
|
top = max(0, int(bounds.y))
|
|
right = min(picture.width, int(round(bounds.right)))
|
|
bottom = min(picture.height, int(round(bounds.bottom)))
|
|
if right - left < 2 or bottom - top < 2:
|
|
return (None, None)
|
|
|
|
pixels = list(picture.crop((left, top, right, bottom)).get_flattened_data())
|
|
threshold = _otsu_threshold(pixels)
|
|
if threshold is None:
|
|
return (None, None)
|
|
|
|
dark = [pixel for pixel in pixels if _luminance(pixel) <= threshold]
|
|
light = [pixel for pixel in pixels if _luminance(pixel) > threshold]
|
|
if not dark or not light:
|
|
return (None, None)
|
|
|
|
# Text strokes normally cover a minority of the box's pixels regardless of
|
|
# whether the text is dark-on-light or light-on-dark, so the smaller of
|
|
# the two luminance clusters is treated as the foreground (text) color.
|
|
foreground, background = (dark, light) if len(dark) <= len(light) else (light, dark)
|
|
return _average_hex(foreground), _average_hex(background)
|
|
|
|
|
|
def _luminance(pixel: tuple[int, int, int]) -> float:
|
|
r, g, b = pixel
|
|
return 0.299 * r + 0.587 * g + 0.114 * b
|
|
|
|
|
|
def _otsu_threshold(pixels: list[tuple[int, int, int]]) -> float | None:
|
|
total = len(pixels)
|
|
if total == 0:
|
|
return None
|
|
|
|
histogram = [0] * 256
|
|
for pixel in pixels:
|
|
histogram[int(_luminance(pixel))] += 1
|
|
|
|
sum_total = sum(level * count for level, count in enumerate(histogram))
|
|
sum_background = 0.0
|
|
weight_background = 0
|
|
best_variance = -1.0
|
|
threshold = 0
|
|
for level, count in enumerate(histogram):
|
|
weight_background += count
|
|
if weight_background == 0:
|
|
continue
|
|
weight_foreground = total - weight_background
|
|
if weight_foreground == 0:
|
|
break
|
|
sum_background += level * count
|
|
mean_background = sum_background / weight_background
|
|
mean_foreground = (sum_total - sum_background) / weight_foreground
|
|
variance = (
|
|
weight_background
|
|
* weight_foreground
|
|
* (mean_background - mean_foreground) ** 2
|
|
)
|
|
if variance > best_variance:
|
|
best_variance = variance
|
|
threshold = level
|
|
return float(threshold)
|
|
|
|
|
|
def _average_hex(pixels: list[tuple[int, int, int]]) -> str:
|
|
count = len(pixels)
|
|
r = sum(pixel[0] for pixel in pixels) // count
|
|
g = sum(pixel[1] for pixel in pixels) // count
|
|
b = sum(pixel[2] for pixel in pixels) // count
|
|
return f"#{r:02x}{g:02x}{b:02x}"
|
|
|
|
|
|
def parse_paddle_result(raw: Any) -> list[OCRBox]:
|
|
boxes: list[OCRBox] = []
|
|
for item in _flatten_pages(raw):
|
|
parsed = _parse_line(item)
|
|
if parsed:
|
|
boxes.append(parsed)
|
|
return boxes
|
|
|
|
|
|
def _image_input(image: bytes | str | Path) -> tuple[str, Path | None]:
|
|
if isinstance(image, bytes):
|
|
handle = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
|
|
try:
|
|
handle.write(image)
|
|
finally:
|
|
handle.close()
|
|
return handle.name, Path(handle.name)
|
|
return os.fspath(image), None
|
|
|
|
|
|
def _flatten_pages(raw: Any) -> Iterable[Any]:
|
|
if raw is None:
|
|
return []
|
|
if isinstance(raw, dict):
|
|
return _dict_lines(raw)
|
|
if isinstance(raw, list):
|
|
flattened: list[Any] = []
|
|
for page in raw:
|
|
if isinstance(page, dict):
|
|
flattened.extend(_dict_lines(page))
|
|
elif _looks_like_ocr_line(page):
|
|
flattened.append(page)
|
|
elif isinstance(page, list):
|
|
flattened.extend(page)
|
|
return flattened
|
|
json_attr = getattr(raw, "json", None)
|
|
if callable(json_attr):
|
|
return _flatten_pages(json_attr())
|
|
if isinstance(json_attr, (dict, list)):
|
|
return _flatten_pages(json_attr)
|
|
return []
|
|
|
|
|
|
def _dict_lines(data: dict[str, Any]) -> list[Any]:
|
|
result = data.get("res")
|
|
payload = result if isinstance(result, dict) else data
|
|
texts = _first_nonempty(payload, "rec_texts", "texts")
|
|
scores = _first_nonempty(payload, "rec_scores", "scores")
|
|
boxes = _first_nonempty(payload, "rec_boxes", "rec_polys", "dt_polys")
|
|
return [
|
|
{
|
|
"text": text,
|
|
"confidence": scores[index] if index < len(scores) else None,
|
|
"points": boxes[index] if index < len(boxes) else None,
|
|
}
|
|
for index, text in enumerate(texts)
|
|
]
|
|
|
|
|
|
def _looks_like_ocr_line(value: Any) -> bool:
|
|
return isinstance(value, (list, tuple)) and len(value) >= 2
|
|
|
|
|
|
def _parse_line(line: Any) -> OCRBox | None:
|
|
if isinstance(line, dict):
|
|
text = line.get("text")
|
|
points = _first_nonempty(line, "points", "box", "bounds")
|
|
confidence = line.get("confidence")
|
|
if not _has_items(text) or not _has_items(points):
|
|
return None
|
|
return OCRBox(
|
|
str(text), _bounds_from_points(points), _float_or_none(confidence)
|
|
)
|
|
|
|
if not _looks_like_ocr_line(line):
|
|
return None
|
|
|
|
points = line[0]
|
|
text_payload = line[1]
|
|
if isinstance(text_payload, (list, tuple)) and text_payload:
|
|
text = text_payload[0]
|
|
confidence = text_payload[1] if len(text_payload) > 1 else None
|
|
else:
|
|
text = text_payload
|
|
confidence = None
|
|
if not text:
|
|
return None
|
|
return OCRBox(str(text), _bounds_from_points(points), _float_or_none(confidence))
|
|
|
|
|
|
def _bounds_from_points(points: Any) -> Bounds:
|
|
if isinstance(points, dict):
|
|
return Bounds.from_dict(points)
|
|
if _has_length(points, 4) and all(isinstance(value, Real) for value in points):
|
|
x1, y1, x2, y2 = [float(value) for value in points]
|
|
return Bounds(x1, y1, x2 - x1, y2 - y1)
|
|
|
|
xs: list[float] = []
|
|
ys: list[float] = []
|
|
for point in points:
|
|
if isinstance(point, dict):
|
|
xs.append(float(point["x"]))
|
|
ys.append(float(point["y"]))
|
|
else:
|
|
xs.append(float(point[0]))
|
|
ys.append(float(point[1]))
|
|
return Bounds(min(xs), min(ys), max(xs) - min(xs), max(ys) - min(ys))
|
|
|
|
|
|
def _float_or_none(value: Any) -> float | None:
|
|
if value is None:
|
|
return None
|
|
return float(value)
|
|
|
|
|
|
def _first_nonempty(data: dict[str, Any], *keys: str) -> Any:
|
|
for key in keys:
|
|
value = data.get(key)
|
|
if _has_items(value):
|
|
return value
|
|
return []
|
|
|
|
|
|
def _has_items(value: Any) -> bool:
|
|
if value is None:
|
|
return False
|
|
try:
|
|
return len(value) > 0
|
|
except TypeError:
|
|
return bool(value)
|
|
|
|
|
|
def _has_length(value: Any, length: int) -> bool:
|
|
try:
|
|
return len(value) == length
|
|
except TypeError:
|
|
return False
|