ba-auto-daily/ba_auto/detector.py

"""Image/color matching helpers (OpenCV-based). Ported from scripts/detect_and_click.py."""
import os

import cv2
import numpy as np
import pytesseract

from ba_auto import config, driver

OCR_SHOT_PATH = os.path.join(config.SCRATCHPAD_DIR, "ocr_live.png")
OCR_UPSCALE = 3

SPARKLE_SHOT_PATH = os.path.join(config.SCRATCHPAD_DIR, "cafe_live.png")
SPARKLE_CLICK_OFFSET = (75, 47)
SPARKLE_THRESHOLD = 0.97

# The cafe camera's zoom level isn't reset before farming, and the sparkle
# icon's on-screen size scales with it (see screenshots/cafe/sparkle/
# 01_sparke_zoomed_centered.png vs 03_sparkle_zoomed_out.png) -- a single
# fixed-scale template match misses whenever the camera isn't at the exact
# zoom the template was captured at. Try a spread of scales instead.
SPARKLE_SCALES = (0.6, 0.7, 0.8, 0.9, 1.0, 1.1, 1.2)


def _masked_template(template):
    b, g, r = cv2.split(template.astype(np.int16))
    yellow_white = ((r > 180) & (g > 140) & (r - b > 60)) | ((r > 200) & (g > 200) & (b > 200))
    mask_plane = (yellow_white.astype(np.uint8)) * 255
    return cv2.merge([mask_plane, mask_plane, mask_plane])


def find_cafe_sparkle():
    driver.screenshot(SPARKLE_SHOT_PATH)
    template_full = cv2.imread(config.CAFE_SPARKLE_TEMPLATE)
    img = cv2.imread(SPARKLE_SHOT_PATH)
    th0, tw0 = template_full.shape[:2]

    best = None
    for scale in SPARKLE_SCALES:
        tw, th = max(1, round(tw0 * scale)), max(1, round(th0 * scale))
        template = cv2.resize(template_full, (tw, th))
        mask = _masked_template(template)

        result = cv2.matchTemplate(img, template, cv2.TM_CCORR_NORMED, mask=mask)
        _, max_val, _, max_loc = cv2.minMaxLoc(result)
        if max_val >= SPARKLE_THRESHOLD and (best is None or max_val > best[0]):
            best = (max_val, max_loc[0], max_loc[1], tw, th, scale)

    if best is None:
        return None

    score, x, y, tw, th, scale = best
    ox, oy = SPARKLE_CLICK_OFFSET
    return (x + tw // 2 + round(ox * scale), y + th // 2 + round(oy * scale), score)


def _color_mask(region, rgb_min, rgb_max):
    x1, y1, x2, y2 = region
    driver.screenshot(OCR_SHOT_PATH)
    img = cv2.imread(OCR_SHOT_PATH)
    crop = img[y1:y2, x1:x2]
    b, g, r = crop[:, :, 0].astype(np.int16), crop[:, :, 1].astype(np.int16), crop[:, :, 2].astype(np.int16)
    (r_lo, g_lo, b_lo), (r_hi, g_hi, b_hi) = rgb_min, rgb_max
    return (r >= r_lo) & (r <= r_hi) & (g >= g_lo) & (g <= g_hi) & (b >= b_lo) & (b <= b_hi)


def region_contains_color(region, rgb_min, rgb_max):
    """Whether any pixel within `region` (x1, y1, x2, y2) falls in the given
    RGB range. Useful for presence checks on small, non-convex glyphs (e.g.
    an arrow chevron) where a single fixed-point probe can land in the
    glyph's own concave gap -- confirmed live: a centroid-based single point
    for story_sweep's region-arrow chevron fell squarely in the notch
    between its two strokes, reading as "absent" even while the arrow was
    clearly rendered a few pixels away. See plan.md Phase 10.
    """
    return bool(_color_mask(region, rgb_min, rgb_max).any())


def find_color_centroid(region, rgb_min, rgb_max):
    """Centroid `(x, y)` of every pixel within `region` (x1, y1, x2, y2)
    falling in the given RGB range, or None if none match. Useful for
    clicking a known-colored button whose exact position varies between
    otherwise-similar dialogs -- e.g. story_sweep's sweep-result screen
    shows the same cyan confirm-button color for its "SKIP" and final "OK"
    states, ~120px apart vertically; finding it by color avoids hardcoding
    both positions. See plan.md Phase 10.
    """
    x1, y1, x2, y2 = region
    mask = _color_mask(region, rgb_min, rgb_max)
    ys, xs = np.nonzero(mask)
    if len(xs) == 0:
        return None
    return (x1 + int(xs.mean()), y1 + int(ys.mean()))


def _ocr_crop(region):
    x1, y1, x2, y2 = region
    driver.screenshot(OCR_SHOT_PATH)
    img = cv2.imread(OCR_SHOT_PATH)
    crop = img[y1:y2, x1:x2]
    gray = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
    # This UI's text is consistently dark-on-light -- a hard black/white
    # threshold measurably fixed real misreads during live calibration
    # (e.g. a stage label's digit misread as a stray extra digit) that
    # persisted across every psm mode until the anti-aliased grey edges were
    # removed. Confirmed live against scratchpad/scroll_up20.png row labels.
    _, thresh = cv2.threshold(gray, 150, 255, cv2.THRESH_BINARY)
    return cv2.resize(thresh, None, fx=OCR_UPSCALE, fy=OCR_UPSCALE, interpolation=cv2.INTER_CUBIC)


def read_text(region, whitelist=None, psm=7, lang="eng"):
    """OCR a pixel rectangle `(x1, y1, x2, y2)` from a fresh screenshot.

    Local equivalent of the reference's core/ocr/ocr.py Baas_ocr client
    (get_region_res) -- without its socket/shared-memory server, which exists
    there to make OCR fast across thousands of automation steps. This project
    only needs occasional single-crop reads (a region number, a stage label),
    so a plain in-process pytesseract call is enough; see CLAUDE.md's "OCR
    policy" and Handoff.md. `lang="eng"` is sufficient for the digit/dash
    labels this project reads (region numbers, "30-1"/"30-A" stage labels) --
    no Japanese trained data is needed for those specific reads.
    """
    crop = _ocr_crop(region)
    tess_config = f"--psm {psm}"
    if whitelist:
        tess_config += f" -c tessedit_char_whitelist={whitelist}"
    return pytesseract.image_to_string(crop, lang=lang, config=tess_config).strip()


def read_int(region, psm=7):
    """Local equivalent of the reference's recognize_int."""
    digits = "".join(ch for ch in read_text(region, whitelist="0123456789", psm=psm) if ch.isdigit())
    return int(digits) if digits else None