宜博 / pixelphone

"""
Local Chinese OCR fallback — PP-OCRv6 tiny (det+rec) via onnxruntime, no cv2.

Why this exists: ui_control_tools.py's find_element/click_by_text reads the
uiautomator accessibility tree, which is strictly more reliable than OCR for
any element exposing text/resource-id/content-desc. OCR is only needed for
the two cases that tree can't cover: icon-only buttons with no accessible
label, and uiautomator returning a null root (AccessibilityService conflict,
see ui_control_tools.py's module docstring).

No opencv: opencv-python has no aarch64-android wheel and fails to build from
source in Termux (broken NDK/cmake toolchain for its `ninja` build dep). Since
phone screenshots are always axis-aligned (no skew), DBNet's polygon output is
collapsed to a plain bounding box — cv2.minAreaRect/warpPerspective are never
needed. Connected components come from scipy.ndimage.label instead of
cv2.findContours; the polygon "unclip" expand step still uses pyclipper
(pure C++ wheel, builds fine, no cv2 dependency) following PaddleOCR's
standard area/perimeter formula.

Models: PP-OCRv6 tiny det+rec ONNX (the edge/IoT tier, 1.5M params) +
ppocrv6_tiny_dict.txt, expected under ocr_models/ next to this file (~6MB,
not committed to git — see .gitignore). ONNX weights come from RapidOCR's
ModelScope mirror (huggingface.co is unreachable from this network); the
dict comes straight from the official PaddlePaddle/PaddleOCR repo, since
RapidOCR's own onnx config doesn't list one yet (see version history below):

    mkdir -p ocr_models && cd ocr_models
    curl -fL -o PP-OCRv6_det_tiny.onnx \
      https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.0/onnx/PP-OCRv6/det/PP-OCRv6_det_tiny.onnx
    curl -fL -o PP-OCRv6_rec_tiny.onnx \
      https://www.modelscope.cn/models/RapidAI/RapidOCR/resolve/v3.9.0/onnx/PP-OCRv6/rec/PP-OCRv6_rec_tiny.onnx
    curl -fL -o ppocrv6_tiny_dict.txt \
      https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/main/ppocr/utils/dict/ppocrv6_tiny_dict.txt

Version history (v4 -> v5 -> v6, all same DBNet-det + CTC-rec pipeline, only
weights/dict swapped each time — verified at each step by introspecting the
onnx I/O shapes and re-running the real-screenshot test, never by assuming):
  - v4: first version found while researching; not a deliberate choice.
  - v5: real accuracy upgrade over v4, fully documented (RapidOCR's
    default_models.yaml carries a dict_url for it). Same input
    normalization, rec height=48, sigmoid-baked det output, CTC
    blank+dict+space decode as v4.
  - v6 (released 2026-06-11, 12 days old at the time of writing): RapidOCR
    already hosts the onnx weights but its default_models.yaml has no
    dict_url for any v6 entry in any backend block (onnxruntime, openvino,
    mnn, paddle) — see github.com/RapidAI/RapidOCR issue #686, an open
    "support pp_ocr_v6" request. The official `paddleocr` pip package can't
    be used here instead: paddleocr -> paddlex[ocr-core] hard-pins
    opencv-contrib-python==4.10.0.84, which hits the exact same
    no-aarch64-android-wheel wall as plain opencv-python. So: took the
    RapidOCR-hosted v6 *tiny* onnx weights, paired them with
    ppocrv6_tiny_dict.txt from the official PaddleOCR GitHub repo (the
    "_tiny" dict, not the 18708-line unified one — confirmed by introspecting
    PP-OCRv6_rec_tiny.onnx's output layer: 6906 classes ==
    6904 dict lines + 1 CTC blank + 1 trailing space), and confirmed det
    output is still a sigmoid-baked [0,1] map by running it on a real
    screenshot rather than assuming. Same pipeline code as v4/v5, no
    postprocessing changes needed.

Runtime deps (Termux native pkg, not pip — opencv-python has no aarch64-android
wheel and pip's fallback source build breaks on Termux's NDK/cmake setup):

    pkg install python-onnxruntime python-numpy python-pillow python-scipy
    pip install pyclipper   # pure C++ ext, builds fine without cv2

and the venv this runs in needs `include-system-site-packages = true` in its
pyvenv.cfg so it can see those pkg-installed packages.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import numpy as np
import onnxruntime as ort
import pyclipper
from PIL import Image
from scipy import ndimage

from ui_control_tools import _do_screenshot, _do_tap, _run

_MODELS_DIR = Path(__file__).resolve().parent / "ocr_models"
_DET_MODEL = _MODELS_DIR / "PP-OCRv6_det_tiny.onnx"
_REC_MODEL = _MODELS_DIR / "PP-OCRv6_rec_tiny.onnx"
_DICT_FILE = _MODELS_DIR / "ppocrv6_tiny_dict.txt"

_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)

_DET_LIMIT_SIDE_LEN = 960
_DET_THRESH = 0.3
_DET_BOX_THRESH = 0.5
_DET_UNCLIP_RATIO = 1.6
_DET_MIN_SIZE = 3

_REC_IMG_H = 48
_REC_MAX_BATCH_WIDTH = 800  # cap per-crop resized width fed to the rec model


class _Engine:
    """Lazily-loaded singleton: model load (~15MB) happens once per process."""

    _det_session: Optional[ort.InferenceSession] = None
    _rec_session: Optional[ort.InferenceSession] = None
    _chars: Optional[list[str]] = None

    @classmethod
    def det(cls) -> ort.InferenceSession:
        if cls._det_session is None:
            cls._det_session = ort.InferenceSession(str(_DET_MODEL),
                                                     providers=["CPUExecutionProvider"])
        return cls._det_session

    @classmethod
    def rec(cls) -> ort.InferenceSession:
        if cls._rec_session is None:
            cls._rec_session = ort.InferenceSession(str(_REC_MODEL),
                                                     providers=["CPUExecutionProvider"])
        return cls._rec_session

    @classmethod
    def chars(cls) -> list[str]:
        if cls._chars is None:
            dict_chars = _DICT_FILE.read_text(encoding="utf-8").splitlines()
            # PaddleOCR CTCLabelDecode convention: index 0 is the CTC blank,
            # the dict file supplies indices 1..N, and a trailing space is
            # appended as the final class.
            cls._chars = ["blank"] + dict_chars + [" "]
        return cls._chars


# ── detection ────────────────────────────────────────────────────────────────

def _det_resize(img: np.ndarray) -> tuple[np.ndarray, float, float]:
    h, w = img.shape[:2]
    ratio = _DET_LIMIT_SIDE_LEN / max(h, w) if max(h, w) > _DET_LIMIT_SIDE_LEN else 1.0
    resize_h = max(32, round(h * ratio / 32) * 32)
    resize_w = max(32, round(w * ratio / 32) * 32)
    resized = Image.fromarray(img).resize((resize_w, resize_h), Image.BILINEAR)
    return np.asarray(resized), resize_w / w, resize_h / h


def _det_preprocess(img: np.ndarray) -> tuple[np.ndarray, float, float]:
    resized, ratio_w, ratio_h = _det_resize(img)
    x = resized.astype(np.float32) / 255.0
    x = (x - _MEAN) / _STD
    x = x.transpose(2, 0, 1)[None].astype(np.float32)  # NCHW
    return x, ratio_w, ratio_h


def _unclip_bbox(x0: int, y0: int, x1: int, y1: int, ratio: float) -> tuple[int, int, int, int]:
    bw, bh = x1 - x0, y1 - y0
    perimeter = 2 * (bw + bh)
    if perimeter == 0:
        return x0, y0, x1, y1
    distance = bw * bh * ratio / perimeter
    rect = [(x0, y0), (x1, y0), (x1, y1), (x0, y1)]
    pco = pyclipper.PyclipperOffset()
    pco.AddPath(rect, pyclipper.JT_ROUND, pyclipper.ET_CLOSEDPOLYGON)
    expanded = pco.Execute(distance)
    pts = np.array(expanded[0])
    return int(pts[:, 0].min()), int(pts[:, 1].min()), int(pts[:, 0].max()), int(pts[:, 1].max())


def detect_boxes(img: np.ndarray) -> list[dict]:
    """Run PP-OCRv4 det on an RGB numpy image, return axis-aligned text boxes.

    Each box is {"bbox": [x0, y0, x1, y1], "score": prob_mean} in the
    original image's pixel coordinates.
    """
    h, w = img.shape[:2]
    x, ratio_w, ratio_h = _det_preprocess(img)
    inp_name = _Engine.det().get_inputs()[0].name
    prob_map = _Engine.det().run(None, {inp_name: x})[0][0, 0]  # (resize_h, resize_w)

    mask = prob_map > _DET_THRESH
    labels, n = ndimage.label(mask)
    boxes = []
    for i in range(1, n + 1):
        ys, xs = np.where(labels == i)
        if ys.size == 0:
            continue
        x0, x1, y0, y1 = xs.min(), xs.max() + 1, ys.min(), ys.max() + 1
        if (x1 - x0) < _DET_MIN_SIZE or (y1 - y0) < _DET_MIN_SIZE:
            continue
        score = float(prob_map[y0:y1, x0:x1].mean())
        if score < _DET_BOX_THRESH:
            continue
        ex0, ey0, ex1, ey1 = _unclip_bbox(int(x0), int(y0), int(x1), int(y1), _DET_UNCLIP_RATIO)
        # map back to original image scale, clamp to bounds
        ox0 = max(0, int(ex0 / ratio_w))
        oy0 = max(0, int(ey0 / ratio_h))
        ox1 = min(w, int(ex1 / ratio_w))
        oy1 = min(h, int(ey1 / ratio_h))
        if ox1 <= ox0 or oy1 <= oy0:
            continue
        boxes.append({"bbox": [ox0, oy0, ox1, oy1], "score": score})
    boxes.sort(key=lambda b: (b["bbox"][1], b["bbox"][0]))
    return boxes


# ── recognition ──────────────────────────────────────────────────────────────

def _rec_preprocess(crop: np.ndarray) -> np.ndarray:
    h, w = crop.shape[:2]
    new_w = min(_REC_MAX_BATCH_WIDTH, max(1, round(w * _REC_IMG_H / h)))
    resized = np.asarray(Image.fromarray(crop).resize((new_w, _REC_IMG_H), Image.BILINEAR))
    x = (resized.astype(np.float32) / 255.0 - 0.5) / 0.5
    x = x.transpose(2, 0, 1)[None].astype(np.float32)
    return x


def _ctc_decode(logits: np.ndarray, chars: list[str]) -> tuple[str, float]:
    idx = logits.argmax(axis=-1)
    probs = logits.max(axis=-1)
    text, confs, prev = [], [], -1
    for t, c in enumerate(idx):
        if c != 0 and c != prev:
            text.append(chars[c] if c < len(chars) else "")
            confs.append(probs[t])
        prev = c
    avg_conf = float(np.mean(confs)) if confs else 0.0
    return "".join(text), avg_conf


def recognize_crop(img: np.ndarray, bbox: list[int]) -> dict:
    """OCR a single bbox crop of img. Returns {"text", "score"}."""
    x0, y0, x1, y1 = bbox
    crop = img[y0:y1, x0:x1]
    x = _rec_preprocess(crop)
    inp_name = _Engine.rec().get_inputs()[0].name
    logits = _Engine.rec().run(None, {inp_name: x})[0][0]  # (T, num_classes)
    text, score = _ctc_decode(logits, _Engine.chars())
    return {"text": text, "score": score}


# ── public entrypoint ────────────────────────────────────────────────────────

def _missing_models() -> list[str]:
    return [str(p) for p in (_DET_MODEL, _REC_MODEL, _DICT_FILE) if not p.exists()]


def ocr_image(path: str) -> dict:
    """Detect + recognize all text on a screenshot. Returns JSON-able dict.

    Each result item: {"text": str, "score": float, "bbox": [x0,y0,x1,y1],
    "center": {"x": int, "y": int}} in the image's own pixel coordinates —
    if the screenshot was taken at full device resolution, these are
    tap()-ready coordinates.
    """
    missing = _missing_models()
    if missing:
        raise FileNotFoundError(f"OCR model files missing: {missing}")
    img = np.asarray(Image.open(path).convert("RGB"))
    boxes = detect_boxes(img)
    results = []
    for b in boxes:
        rec = recognize_crop(img, b["bbox"])
        if not rec["text"]:
            continue
        x0, y0, x1, y1 = b["bbox"]
        results.append({
            "text": rec["text"],
            "score": round(min(rec["score"], b["score"]), 4),
            "bbox": b["bbox"],
            "center": {"x": (x0 + x1) // 2, "y": (y0 + y1) // 2},
        })
    return {"ok": True, "path": path, "count": len(results), "results": results}


# ── MCP tool fallbacks for find_element / click_by_text ────────────────────

def _do_ocr_find_text(text: str = "", partial: bool = True, limit: int = 20,
                       path: str = "") -> dict:
    owns_file = not path
    if not path:
        shot = _do_screenshot()
        if not shot.get("ok"):
            return {"ok": False, "error": shot.get("error", "screenshot failed")}
        path = shot["path"]
    try:
        result = ocr_image(path)
    except FileNotFoundError as e:
        return {"ok": False, "error": str(e)}
    finally:
        if owns_file:
            _run(f"rm -f {path}")

    matches = []
    for item in result["results"]:
        hay = item["text"]
        ok = (text in hay) if partial else (text == hay)
        if text and not ok:
            continue
        matches.append(item)
    matches = matches[:max(1, limit)]
    return {"ok": bool(matches), "query": {"text": text, "partial": partial},
            "count": len(matches), "matches": matches}


def _do_ocr_click_text(text: str, partial: bool = True, index: int = 0) -> dict:
    found = _do_ocr_find_text(text=text, partial=partial, limit=max(index + 1, 20))
    if not found.get("ok"):
        return {"ok": False, "error": found.get("error") or f"no text matched {text!r} via OCR"}
    matches = found["matches"]
    if index < 0 or index >= len(matches):
        index = 0
    target = matches[index]
    c = target["center"]
    tap = _do_tap(c["x"], c["y"])
    return {"ok": tap.get("ok", False), "action": "ocr_click_text", "text": text,
            "tapped": c, "match": target, "matched_total": len(matches)}


def register_ocr_tools(mcp) -> None:
    """Register OCR-fallback tools onto an existing FastMCP instance."""

    @mcp.tool()
    def ocr_find_text(text: str = "", partial: bool = True, limit: int = 20,
                       path: str = "") -> str:
        """Find on-screen text via local OCR — fallback for find_element.

        Use only when find_element can't see the element: uiautomator
        returning a null root (AccessibilityService conflict), or a label
        with no resource-id/content-desc for find_element to match. This is
        whole-screen PP-OCRv5 inference (~1-7s on-device) reading rendered
        pixels, not the accessibility tree — try find_element first. Boxes
        are per detected text line/cluster, not per substring: a partial
        match inside a long sentence returns that whole line's bounds, not
        just the matched word's. Icon-only buttons with no visible text are
        never found this way.

        Args:
            text: substring (or exact, see partial) to match against
                  recognized text. Empty returns every text region found.
            partial: substring match (default True) vs exact match.
            limit: max matches to return (default 20).
            path: OCR an existing on-device screenshot instead of capturing a
                  fresh one (a fresh capture is deleted after OCR; a path you
                  pass in is left alone).
        """
        return json.dumps(_do_ocr_find_text(text, partial, limit, path), ensure_ascii=False)

    @mcp.tool()
    def ocr_click_text(text: str, partial: bool = True, index: int = 0) -> str:
        """Find visible text via local OCR and tap its center — fallback for click_by_text.

        Use only when click_by_text fails (uiautomator null root, or the
        label has no resource-id/content-desc for find_element to match).
        Do not use on payment, deletion, or sensitive-setting screens.

        Args:
            text: visible text to locate via OCR and tap.
            partial: substring match (default True) vs exact match.
            index: which match to tap if several share the text (default 0).
        """
        return json.dumps(_do_ocr_click_text(text, partial, index), ensure_ascii=False)


if __name__ == "__main__":
    import sys
    print(json.dumps(ocr_image(sys.argv[1]), ensure_ascii=False, indent=2))