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


def read_int_on_heart_badge(region, psm=7):
    """OCR a small dark-navy digit rendered on lesson.py's pink/magenta
    heart-shaped affection badge.

    read_int's plain grayscale threshold (tuned for this UI's normal
    dark-text-on-light-card look) misreads these: the heart's own outline
    stroke is a darker, saturated magenta whose grayscale value happens to
    fall on the same side of the threshold as the digit glyph, so it
    survives as stray black marks tesseract sometimes fuses into extra
    digits -- confirmed live, "13" -> "113", "19" -> "119". The digit color
    is reliably R < G (navy/blue-grey) while every pink/magenta badge tone
    sampled (fill and outline, light and dark) is R > G, so masking on that
    channel relationship instead of raw brightness cleanly drops the badge
    shape and keeps just the glyph.
    """
    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)
    ink = (r < g) & (np.maximum(np.maximum(r, g), b) < 170)
    binary = np.where(ink, 0, 255).astype(np.uint8)
    upscaled = cv2.resize(binary, None, fx=OCR_UPSCALE, fy=OCR_UPSCALE, interpolation=cv2.INTER_CUBIC)
    tess_config = f"--psm {psm} -c tessedit_char_whitelist=0123456789"
    text = pytesseract.image_to_string(upscaled, lang="eng", config=tess_config).strip()
    digits = "".join(ch for ch in text if ch.isdigit())
    return int(digits) if digits else None