diff --git a/core/models.py b/core/models.py index 307cee3..df07189 100644 --- a/core/models.py +++ b/core/models.py @@ -72,6 +72,7 @@ class Device: _STATE_FIELDS = ("enabled", "clickable", "selected", "checked", "focused") +_COLOR_FIELDS = ("foreground_color", "background_color") @dataclass @@ -89,6 +90,11 @@ class SceneElement: selected: bool | None = None checked: bool | None = None focused: bool | None = None + # Text/background color sampled from the screenshot pixels under the OCR + # box ("#rrggbb"). None means unavailable (not OCR-sourced, or sampling + # failed), not "no color". + foreground_color: str | None = None + background_color: str | None = None @property def center(self) -> tuple[float, float]: @@ -104,7 +110,7 @@ class SceneElement: } if self.source: data["source"] = self.source - for field_name in _STATE_FIELDS: + for field_name in _STATE_FIELDS + _COLOR_FIELDS: value = getattr(self, field_name) if value is not None: data[field_name] = value @@ -124,6 +130,8 @@ class SceneElement: selected=data.get("selected"), checked=data.get("checked"), focused=data.get("focused"), + foreground_color=data.get("foreground_color"), + background_color=data.get("background_color"), ) diff --git a/perception/ocr.py b/perception/ocr.py index 27f9ce8..afc091f 100644 --- a/perception/ocr.py +++ b/perception/ocr.py @@ -4,7 +4,7 @@ import os import tempfile import logging from collections.abc import Iterable -from dataclasses import dataclass +from dataclasses import dataclass, replace from numbers import Real from pathlib import Path from typing import Any @@ -19,6 +19,8 @@ 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( @@ -28,6 +30,8 @@ class OCRBox: bounds=self.bounds, confidence=self.confidence, source="ocr", + foreground_color=self.foreground_color, + background_color=self.background_color, ) @@ -54,7 +58,8 @@ class PaddleOCREngine: raw = engine.predict(input=image_input) else: raw = engine.ocr(image_input, cls=True) - return parse_paddle_result(raw) + boxes = parse_paddle_result(raw) + return _with_sampled_colors(boxes, image_input) finally: if temp_path: temp_path.unlink(missing_ok=True) @@ -83,6 +88,106 @@ def run_ocr( 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): diff --git a/pyproject.toml b/pyproject.toml index 3870e3d..4a51bc4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "openai>=1.0.0", "paddlepaddle>=3.0.0,<3.3.0", "paddleocr>=3.0.0", + "pillow>=12.0.0", ] [build-system] diff --git a/runtime/planner_prompts.py b/runtime/planner_prompts.py index 8e1fda1..b465d2d 100644 --- a/runtime/planner_prompts.py +++ b/runtime/planner_prompts.py @@ -19,6 +19,14 @@ interacted with (do not tap it); `selected`/`checked`/`focused` describe its current toggle/focus state and are useful for deciding whether an action is already done or still needed. +Text elements sourced from OCR may also carry `foreground_color` and +`background_color` ("#rrggbb", sampled from the screenshot pixels under that +text). These are omitted when the element is not OCR-sourced or sampling +failed — omitted does NOT mean "no color", treat it as unknown. Use them only +as a secondary signal (e.g. to tell an active/highlighted item apart from an +inactive one with the same text) and prefer bounds/text/screenshot evidence +when they disagree. + Before calling a tool, output a short text block (1-2 sentences): 1. If this is the first step, state what you intend to do and why. 2. Otherwise, first assess whether the previous action achieved its intended diff --git a/tests/test_ocr.py b/tests/test_ocr.py index d3ed929..9d0b1f0 100644 --- a/tests/test_ocr.py +++ b/tests/test_ocr.py @@ -2,9 +2,16 @@ from __future__ import annotations import numpy as np import pytest +from PIL import Image -from core.models import Bounds -from perception.ocr import OCRBox, PaddleOCREngine, parse_paddle_result, run_ocr +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: @@ -77,3 +84,64 @@ def test_run_ocr_degrades_when_ocr_engine_raises() -> None: 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 diff --git a/uv.lock b/uv.lock index d4e2796..43221a1 100644 --- a/uv.lock +++ b/uv.lock @@ -406,6 +406,7 @@ dependencies = [ { name = "openai" }, { name = "paddleocr" }, { name = "paddlepaddle" }, + { name = "pillow" }, ] [package.dev-dependencies] @@ -422,6 +423,7 @@ requires-dist = [ { name = "openai", specifier = ">=1.0.0" }, { name = "paddleocr", specifier = ">=3.0.0" }, { name = "paddlepaddle", specifier = ">=3.0.0,<3.3.0" }, + { name = "pillow", specifier = ">=12.0.0" }, ] [package.metadata.requires-dev]