96 lines
3.0 KiB
Python
96 lines
3.0 KiB
Python
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,
|
|
ocr_elements=list(ocr_elements),
|
|
)
|
|
|
|
|
|
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)
|