sos-bot/core/vision.py

"""
Vision — in-memory screenshot capture + OpenCV template matching.
No files are written to disk for screenshots; everything lives in numpy arrays.
"""
import logging
import numpy as np
import cv2

logger = logging.getLogger("sos.vision")


class Vision:
    """Captures screenshots and provides template-matching utilities."""

    def __init__(self, adb_client):
        self.adb = adb_client
        self._screenshot = None       # current frame (BGR numpy array)
        self._screenshot_gray = None  # grayscale version (cached)
        self._size = None             # (width, height)

    # ── capture ──────────────────────────────────────────────

    def capture(self) -> np.ndarray:
        """
        Take a fresh screenshot from the device.
        Returns BGR numpy array.  Never touches disk.
        """
        png_bytes = self.adb.screencap()
        # Decode PNG from memory buffer
        arr = np.frombuffer(png_bytes, dtype=np.uint8)
        img = cv2.imdecode(arr, cv2.IMREAD_COLOR)
        if img is None:
            raise RuntimeError("Failed to decode screencap — corrupt PNG data?")
        self._screenshot = img
        self._screenshot_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        self._size = (img.shape[1], img.shape[0])
        return img

    @property
    def current(self) -> np.ndarray:
        """Return the most recent screenshot (call capture() first)."""
        if self._screenshot is None:
            raise RuntimeError("No screenshot — call capture() first")
        return self._screenshot

    @property
    def size(self):
        """(width, height) of the device screen."""
        if self._size is None:
            self.capture()
        return self._size

    # ── template matching ────────────────────────────────────

    def find_template(
        self,
        template_path: str,
        confidence: float = 0.8,
        region: tuple = None,
    ) -> tuple | None:
        """
        Locate a template image on the current screenshot.

        Args:
            template_path: Path to the reference PNG.
            confidence:    Minimum match ratio (0-1).  Higher = stricter.
            region:        (x1, y1, x2, y2) subregion to search, or None = full.

        Returns:
            (center_x, center_y) of the best match, or None.
        """
        template_rgb = cv2.imread(template_path, cv2.IMREAD_COLOR)
        if template_rgb is None:
            logger.warning("Cannot load template: %s", template_path)
            return None

        template_gray = cv2.cvtColor(template_rgb, cv2.COLOR_BGR2GRAY)
        tw, th = template_gray.shape[1], template_gray.shape[0]

        # Use the cached grayscale screenshot
        search_img = self._screenshot_gray

        # Crop to region if specified
        if region:
            x1, y1, x2, y2 = region
            search_img = search_img[y1:y2, x1:x2]

        if search_img.shape[0] < th or search_img.shape[1] < tw:
            logger.debug("Template %s bigger than search region", template_path)
            return None

        result = cv2.matchTemplate(search_img, template_gray, cv2.TM_CCOEFF_NORMED)
        _, max_val, _, max_loc = cv2.minMaxLoc(result)

        if max_val < confidence:
            logger.debug(
                "Template %s: best=%.3f < threshold=%.3f",
                template_path, max_val, confidence,
            )
            return None

        cx = max_loc[0] + tw // 2
        cy = max_loc[1] + th // 2
        if region:
            cx += region[0]
            cy += region[1]

        logger.info(
            "Found %s at (%.3f conf) → click (%d, %d)",
            template_path, max_val, cx, cy,
        )
        return (cx, cy)

    def find_multiple(
        self,
        template_path: str,
        confidence: float = 0.8,
        max_results: int = 5,
        region: tuple = None,
    ) -> list[tuple]:
        """
        Find ALL matches of a template (non-max-suppressed).

        Returns list of (cx, cy) centers sorted by confidence descending.
        """
        template_rgb = cv2.imread(template_path, cv2.IMREAD_COLOR)
        if template_rgb is None:
            return []
        template_gray = cv2.cvtColor(template_rgb, cv2.COLOR_BGR2GRAY)
        tw, th = template_gray.shape[1], template_gray.shape[0]

        search_img = self._screenshot_gray
        if region:
            x1, y1, x2, y2 = region
            search_img = search_img[y1:y2, x1:x2]

        result = cv2.matchTemplate(search_img, template_gray, cv2.TM_CCOEFF_NORMED)
        locations = np.where(result >= confidence)
        points = []
        for pt in zip(*locations[::-1]):  # (x, y) pairs
            cx = pt[0] + tw // 2
            cy = pt[1] + th // 2
            if region:
                cx += region[0]
                cy += region[1]
            points.append((cx, cy))

        # Non-max suppression: group nearby detections
        if not points:
            return []
        points = np.array(points)
        indices = cv2.dnn.NMSBoxes(
            [(*p, tw, th) for p in points],
            [1.0] * len(points),
            0.5,
            0.3,
        )
        if len(indices) == 0:
            return []
        result_pts = [tuple(points[i]) for i in indices.flatten()]
        return result_pts[:max_results]

    # ── pixel checking ───────────────────────────────────────

    def pixel_color(self, x: int, y: int) -> tuple:
        """Get the BGR color of a single pixel on the current screenshot."""
        return tuple(self.current[y, x].tolist())

    def color_match(self, x: int, y: int, target_bgr: tuple, tolerance: int = 10) -> bool:
        """Check if pixel at (x,y) is within tolerance of target_bgr."""
        actual = self.pixel_color(x, y)
        return all(abs(actual[c] - target_bgr[c]) <= tolerance for c in range(3))

    # ── debug ────────────────────────────────────────────────

    def save_debug_screenshot(self, path: str):
        """Save current screenshot to disk for debugging (only when requested)."""
        if self._screenshot is not None:
            cv2.imwrite(path, self._screenshot)
            logger.info("Debug screenshot saved: %s", path)