feat: checkpoint device agent runtime milestones
This commit is contained in:
@@ -0,0 +1,2 @@
|
||||
"""Screen perception helpers for OCR, UI tree parsing, and Scene fusion."""
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from core.models import Scene
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class IconSearchResult:
|
||||
found: bool
|
||||
name: str
|
||||
x: float | None = None
|
||||
y: float | None = None
|
||||
reason: str | None = None
|
||||
|
||||
def to_dict(self) -> dict[str, object]:
|
||||
return {
|
||||
"found": self.found,
|
||||
"name": self.name,
|
||||
"x": self.x,
|
||||
"y": self.y,
|
||||
"reason": self.reason,
|
||||
}
|
||||
|
||||
|
||||
def find_icon(scene: Scene, name: str) -> IconSearchResult:
|
||||
query = name.casefold()
|
||||
for element in scene.elements:
|
||||
if element.type not in {"image", "icon", "button"}:
|
||||
continue
|
||||
label = (element.text or element.id).casefold()
|
||||
if query in label:
|
||||
x, y = element.center
|
||||
return IconSearchResult(found=True, name=name, x=x, y=y)
|
||||
return IconSearchResult(found=False, name=name, reason="not found")
|
||||
|
||||
@@ -0,0 +1,188 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from collections.abc import Iterable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from core.models import Bounds, SceneElement
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OCRBox:
|
||||
text: str
|
||||
bounds: Bounds
|
||||
confidence: float | 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",
|
||||
)
|
||||
|
||||
|
||||
class PaddleOCREngine:
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
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)
|
||||
return parse_paddle_result(raw)
|
||||
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 ImportError:
|
||||
if strict:
|
||||
raise
|
||||
return []
|
||||
return [box.to_scene_element(f"ocr-{index:03d}") for index, box in enumerate(boxes)]
|
||||
|
||||
|
||||
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)
|
||||
return []
|
||||
|
||||
|
||||
def _dict_lines(data: dict[str, Any]) -> list[Any]:
|
||||
payload = data.get("res") if isinstance(data.get("res"), dict) else data
|
||||
texts = payload.get("rec_texts") or payload.get("texts") or []
|
||||
scores = payload.get("rec_scores") or payload.get("scores") or []
|
||||
boxes = (
|
||||
payload.get("rec_boxes")
|
||||
or payload.get("rec_polys")
|
||||
or payload.get("dt_polys")
|
||||
or []
|
||||
)
|
||||
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 = line.get("points") or line.get("box") or line.get("bounds")
|
||||
confidence = line.get("confidence")
|
||||
if not text or not 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 (
|
||||
isinstance(points, (list, tuple))
|
||||
and len(points) == 4
|
||||
and all(isinstance(value, (int, float)) 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)
|
||||
|
||||
@@ -0,0 +1,36 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
from core.models import Scene
|
||||
from perception.ocr import PaddleOCREngine, run_ocr
|
||||
from perception.scene_builder import build_scene as build_fused_scene
|
||||
from perception.scene_builder import infer_png_size
|
||||
from perception.ui_parser import parse_ui_tree
|
||||
|
||||
|
||||
class PerceptionProvider(ABC):
|
||||
@abstractmethod
|
||||
def build_scene(self, screenshot: bytes, tree: Any) -> Scene:
|
||||
"""Build a normalized scene from raw device observations."""
|
||||
|
||||
|
||||
class DefaultPerceptionProvider(PerceptionProvider):
|
||||
def __init__(self, *, ocr_engine: PaddleOCREngine | None = None) -> None:
|
||||
self._ocr_engine = ocr_engine
|
||||
|
||||
def build_scene(self, screenshot: bytes, tree: Any) -> Scene:
|
||||
width, height = infer_png_size(screenshot)
|
||||
return build_fused_scene(
|
||||
screen_width=width,
|
||||
screen_height=height,
|
||||
ui_elements=parse_ui_tree(tree),
|
||||
ocr_elements=run_ocr(screenshot, engine=self._ocr_engine),
|
||||
)
|
||||
|
||||
|
||||
class NullPerceptionProvider(PerceptionProvider):
|
||||
def build_scene(self, screenshot: bytes, tree: Any) -> Scene:
|
||||
width, height = infer_png_size(screenshot)
|
||||
return Scene(width=width, height=height, elements=[])
|
||||
@@ -0,0 +1,91 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import replace
|
||||
from pathlib import Path
|
||||
from struct import unpack
|
||||
|
||||
from core.models import Bounds, Scene, SceneElement
|
||||
|
||||
|
||||
def build_scene(
|
||||
*,
|
||||
screen_width: int,
|
||||
screen_height: int,
|
||||
ui_elements: list[SceneElement] | None = None,
|
||||
ocr_elements: list[SceneElement] | None = None,
|
||||
iou_threshold: float = 0.5,
|
||||
) -> Scene:
|
||||
ui_elements = ui_elements or []
|
||||
ocr_elements = ocr_elements or []
|
||||
merged: list[SceneElement] = []
|
||||
used_ocr: set[int] = set()
|
||||
|
||||
for ui_index, ui_element in enumerate(ui_elements):
|
||||
best_index: int | None = None
|
||||
best_iou = 0.0
|
||||
for ocr_index, ocr_element in enumerate(ocr_elements):
|
||||
if ocr_index in used_ocr:
|
||||
continue
|
||||
score = bbox_iou(ui_element.bounds, ocr_element.bounds)
|
||||
if score > best_iou:
|
||||
best_iou = score
|
||||
best_index = ocr_index
|
||||
|
||||
element = _with_id(ui_element, f"ui-{ui_index:03d}")
|
||||
if best_index is not None and best_iou >= iou_threshold:
|
||||
used_ocr.add(best_index)
|
||||
ocr_element = ocr_elements[best_index]
|
||||
element = replace(
|
||||
element,
|
||||
text=element.text or ocr_element.text,
|
||||
confidence=_best_confidence(element.confidence, ocr_element.confidence),
|
||||
)
|
||||
merged.append(element)
|
||||
|
||||
for ocr_index, ocr_element in enumerate(ocr_elements):
|
||||
if ocr_index in used_ocr:
|
||||
continue
|
||||
merged.append(_with_id(ocr_element, f"ocr-{ocr_index:03d}"))
|
||||
|
||||
return Scene(width=screen_width, height=screen_height, elements=merged)
|
||||
|
||||
|
||||
def bbox_iou(first: Bounds, second: Bounds) -> float:
|
||||
x_left = max(first.x, second.x)
|
||||
y_top = max(first.y, second.y)
|
||||
x_right = min(first.right, second.right)
|
||||
y_bottom = min(first.bottom, second.bottom)
|
||||
if x_right <= x_left or y_bottom <= y_top:
|
||||
return 0.0
|
||||
|
||||
intersection = (x_right - x_left) * (y_bottom - y_top)
|
||||
first_area = first.width * first.height
|
||||
second_area = second.width * second.height
|
||||
union = first_area + second_area - intersection
|
||||
if union <= 0:
|
||||
return 0.0
|
||||
return intersection / union
|
||||
|
||||
|
||||
def infer_png_size(image: bytes | str | Path | None) -> tuple[int, int]:
|
||||
if image is None:
|
||||
return (0, 0)
|
||||
data = Path(image).read_bytes() if not isinstance(image, bytes) else image
|
||||
if len(data) >= 24 and data[:8] == b"\x89PNG\r\n\x1a\n":
|
||||
width, height = unpack(">II", data[16:24])
|
||||
return (int(width), int(height))
|
||||
return (0, 0)
|
||||
|
||||
|
||||
def _with_id(element: SceneElement, fallback_id: str) -> SceneElement:
|
||||
if element.id:
|
||||
return element
|
||||
return replace(element, id=fallback_id)
|
||||
|
||||
|
||||
def _best_confidence(first: float | None, second: float | None) -> float | None:
|
||||
values = [value for value in (first, second) if value is not None]
|
||||
if not values:
|
||||
return None
|
||||
return max(values)
|
||||
|
||||
@@ -0,0 +1,151 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import xml.etree.ElementTree as ET
|
||||
from typing import Any
|
||||
|
||||
from core.models import Bounds, SceneElement
|
||||
|
||||
_ANDROID_BOUNDS = re.compile(r"\[(?P<x1>-?\d+),(?P<y1>-?\d+)\]\[(?P<x2>-?\d+),(?P<y2>-?\d+)\]")
|
||||
|
||||
|
||||
def parse_ui_tree(raw_tree: Any) -> list[SceneElement]:
|
||||
if raw_tree is None:
|
||||
return []
|
||||
if isinstance(raw_tree, str):
|
||||
return _parse_xml(raw_tree)
|
||||
if isinstance(raw_tree, dict):
|
||||
elements: list[SceneElement] = []
|
||||
_parse_dict_node(raw_tree, elements)
|
||||
return _with_stable_ids(elements, "ui")
|
||||
if isinstance(raw_tree, list):
|
||||
elements = []
|
||||
for node in raw_tree:
|
||||
if isinstance(node, dict):
|
||||
_parse_dict_node(node, elements)
|
||||
return _with_stable_ids(elements, "ui")
|
||||
return []
|
||||
|
||||
|
||||
def _parse_xml(raw_tree: str) -> list[SceneElement]:
|
||||
if not raw_tree.strip():
|
||||
return []
|
||||
root = ET.fromstring(raw_tree)
|
||||
elements = []
|
||||
for node in root.iter():
|
||||
bounds = _bounds_from_attrs(node.attrib)
|
||||
if bounds is None or bounds.width <= 0 or bounds.height <= 0:
|
||||
continue
|
||||
elements.append(
|
||||
SceneElement(
|
||||
id="",
|
||||
type=_normalize_type(_local_name(node.tag), node.attrib),
|
||||
text=_text_from_attrs(node.attrib),
|
||||
bounds=bounds,
|
||||
confidence=1.0,
|
||||
source="ui",
|
||||
)
|
||||
)
|
||||
return _with_stable_ids(elements, "ui")
|
||||
|
||||
|
||||
def _parse_dict_node(node: dict[str, Any], elements: list[SceneElement]) -> None:
|
||||
attrs = dict(node)
|
||||
children = attrs.pop("children", None) or attrs.pop("nodes", None) or []
|
||||
bounds = _bounds_from_attrs(attrs)
|
||||
if bounds and bounds.width > 0 and bounds.height > 0:
|
||||
elements.append(
|
||||
SceneElement(
|
||||
id="",
|
||||
type=_normalize_type(
|
||||
str(attrs.get("type") or attrs.get("class") or "unknown"),
|
||||
attrs,
|
||||
),
|
||||
text=_text_from_attrs(attrs),
|
||||
bounds=bounds,
|
||||
confidence=1.0,
|
||||
source="ui",
|
||||
)
|
||||
)
|
||||
for child in children:
|
||||
if isinstance(child, dict):
|
||||
_parse_dict_node(child, elements)
|
||||
|
||||
|
||||
def _with_stable_ids(elements: list[SceneElement], prefix: str) -> list[SceneElement]:
|
||||
return [
|
||||
SceneElement(
|
||||
id=element.id or f"{prefix}-{index:03d}",
|
||||
type=element.type,
|
||||
text=element.text,
|
||||
bounds=element.bounds,
|
||||
confidence=element.confidence,
|
||||
source=element.source,
|
||||
)
|
||||
for index, element in enumerate(elements)
|
||||
]
|
||||
|
||||
|
||||
def _bounds_from_attrs(attrs: dict[str, Any]) -> Bounds | None:
|
||||
if all(key in attrs for key in ("x", "y", "width", "height")):
|
||||
return Bounds(
|
||||
float(attrs["x"]),
|
||||
float(attrs["y"]),
|
||||
float(attrs["width"]),
|
||||
float(attrs["height"]),
|
||||
)
|
||||
if all(key in attrs for key in ("left", "top", "right", "bottom")):
|
||||
left = float(attrs["left"])
|
||||
top = float(attrs["top"])
|
||||
right = float(attrs["right"])
|
||||
bottom = float(attrs["bottom"])
|
||||
return Bounds(left, top, right - left, bottom - top)
|
||||
raw_bounds = attrs.get("bounds") or attrs.get("rect")
|
||||
if isinstance(raw_bounds, dict):
|
||||
return _bounds_from_attrs(raw_bounds)
|
||||
if isinstance(raw_bounds, str):
|
||||
match = _ANDROID_BOUNDS.fullmatch(raw_bounds.strip())
|
||||
if match:
|
||||
x1 = float(match.group("x1"))
|
||||
y1 = float(match.group("y1"))
|
||||
x2 = float(match.group("x2"))
|
||||
y2 = float(match.group("y2"))
|
||||
return Bounds(x1, y1, x2 - x1, y2 - y1)
|
||||
return None
|
||||
|
||||
|
||||
def _text_from_attrs(attrs: dict[str, Any]) -> str | None:
|
||||
for key in ("label", "name", "text", "value", "placeholder"):
|
||||
value = attrs.get(key)
|
||||
if value not in (None, ""):
|
||||
return str(value)
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_type(raw_type: str, attrs: dict[str, Any]) -> str:
|
||||
candidate = (
|
||||
attrs.get("role")
|
||||
or attrs.get("type")
|
||||
or attrs.get("class")
|
||||
or attrs.get("className")
|
||||
or raw_type
|
||||
)
|
||||
normalized = str(candidate).split(".")[-1].replace("XCUIElementType", "").lower()
|
||||
if "button" in normalized:
|
||||
return "button"
|
||||
if "textfield" in normalized or "textarea" in normalized or "input" in normalized:
|
||||
return "input"
|
||||
if "statictext" in normalized or normalized in {"text", "label"}:
|
||||
return "text"
|
||||
if "image" in normalized or "icon" in normalized:
|
||||
return "image"
|
||||
if "cell" in normalized or "row" in normalized:
|
||||
return "cell"
|
||||
if "window" in normalized:
|
||||
return "window"
|
||||
return normalized or "unknown"
|
||||
|
||||
|
||||
def _local_name(tag: str) -> str:
|
||||
return tag.rsplit("}", 1)[-1]
|
||||
|
||||
Reference in New Issue
Block a user