grid_classifier.py

"""
ONNX Grid Vision Classifier
============================
Server-side ONNX inference engine for grid-based image classification research.

Supports two inference modes:
  - 3x3 grids: per-class binary classifiers (224x224 input, ImageNet normalization)
  - 4x4 grids: YOLOv5 instance segmentation (320x320 input)

Recognizable object categories (15 classes):
  bicycle, bus, car, chimney, crosswalk, fire_hydrant, motorcycle,
  mountain_or_hill, palm_tree, parking_meter, stair, taxi, tractor,
  traffic_light, bridge

Per-class classifier models exist for: bicycle, bus, car, crosswalk,
fire_hydrant, motorcycle, traffic_light (3x3 mode).
YOLOv5-seg handles all 15 classes (4x4 mode).
"""

import os
import logging
from typing import List, Optional

import numpy as np
from PIL import Image
import onnxruntime as ort

logger = logging.getLogger(__name__)

# ── Constants ────────────────────────────────────────────────────────────────

# YOLOv5-seg class names (must match training order)
YOLO_CLASSES = [
    "bicycle", "bus", "car", "chimney", "crosswalk",
    "fire_hydrant", "motorcycle", "mountain_or_hill",
    "palm_tree", "parking_meter", "stair", "taxi",
    "tractor", "traffic_light", "bridge",
]

# Task label → normalized internal label mapping
LABEL_MAP = {
    "bicycles": "bicycle",
    "bicycle": "bicycle",
    "bridges": "bridge",
    "bridge": "bridge",
    "buses": "bus",
    "bus": "bus",
    "cars": "car",
    "car": "car",
    "chimneys": "chimney",
    "chimney": "chimney",
    "crosswalks": "crosswalk",
    "crosswalk": "crosswalk",
    "fire hydrants": "fire_hydrant",
    "fire hydrant": "fire_hydrant",
    "fire_hydrants": "fire_hydrant",
    "motorcycles": "motorcycle",
    "motorcycle": "motorcycle",
    "mountains": "mountain_or_hill",
    "mountain": "mountain_or_hill",
    "palm trees": "palm_tree",
    "palm tree": "palm_tree",
    "parking meters": "parking_meter",
    "parking meter": "parking_meter",
    "stairs": "stair",
    "stair": "stair",
    "taxis": "taxi",
    "taxi": "taxi",
    "tractors": "tractor",
    "tractor": "tractor",
    "traffic lights": "traffic_light",
    "traffic light": "traffic_light",
    "vehicles": "car",
    "vehicle": "car",
}

# ImageNet mean and std for classifier preprocessing
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)


class GridClassifier:
    """ONNX Runtime inference engine for grid-based image classification."""

    def __init__(self, models_dir: str = "models"):
        self.models_dir = models_dir
        self.loaded = False

        # Segmentation models (4x4 grid)
        self.seg_model = None   # yolov5-seg.ort
        self.nms_model = None   # nms-yolov5-det.ort
        self.mask_model = None  # mask-yolov5-seg.ort

        # Classification models (3x3 grid)
        self.classifiers: dict = {}

    # ── Model Loading ────────────────────────────────────────────────────────

    def load_models(self):
        """Load all ONNX models into memory."""
        sess_opts = ort.SessionOptions()
        sess_opts.inter_op_num_threads = os.cpu_count() or 4
        sess_opts.intra_op_num_threads = os.cpu_count() or 4
        sess_opts.execution_mode = ort.ExecutionMode.ORT_PARALLEL

        providers = ["CPUExecutionProvider"]

        # Load segmentation models for 4x4
        seg_path = os.path.join(self.models_dir, "yolov5-seg.ort")
        nms_path = os.path.join(self.models_dir, "nms-yolov5-det.ort")
        mask_path = os.path.join(self.models_dir, "mask-yolov5-seg.ort")

        if os.path.exists(seg_path):
            self.seg_model = ort.InferenceSession(seg_path, sess_opts, providers=providers)
            logger.info(f"Loaded segmentation model: {seg_path}")
        else:
            logger.warning(f"Segmentation model not found: {seg_path}")

        if os.path.exists(nms_path):
            self.nms_model = ort.InferenceSession(nms_path, sess_opts, providers=providers)
            logger.info(f"Loaded NMS model: {nms_path}")
        else:
            logger.warning(f"NMS model not found: {nms_path}")

        if os.path.exists(mask_path):
            self.mask_model = ort.InferenceSession(mask_path, sess_opts, providers=providers)
            logger.info(f"Loaded mask model: {mask_path}")
        else:
            logger.warning(f"Mask model not found: {mask_path}")

        # Load per-class classifiers for 3x3
        classifier_labels = [
            "bicycle", "bus", "car", "crosswalk",
            "fire_hydrant", "motorcycle", "traffic_light",
        ]
        for label in classifier_labels:
            path = os.path.join(self.models_dir, f"{label}.ort")
            if os.path.exists(path):
                self.classifiers[label] = ort.InferenceSession(
                    path, sess_opts, providers=providers
                )
                logger.info(f"Loaded classifier: {path}")
            else:
                logger.warning(f"Classifier not found: {path}")

        self.loaded = True
        logger.info(
            f"ONNX models loaded: {len(self.classifiers)} classifiers, "
            f"segmentation={'yes' if self.seg_model else 'no'}"
        )

    # ── Label Normalization ──────────────────────────────────────────────────

    @staticmethod
    def normalize_label(raw_text: str) -> Optional[str]:
        """
        Extract and normalize a task label from instruction text.

        Input examples:
            "Select all squares with a crosswalk"
            "Select all images with a bus"
            "Select all images containing a fire hydrant"
        """
        text = raw_text.strip().lower()

        # Strip common prefixes
        for prefix in [
            "select all squares with ",
            "select all images with ",
            "select all squares containing ",
            "select all images containing ",
            "select all images with ",
            "tap all squares with ",
            "click all squares with ",
            "select all ",
        ]:
            if text.startswith(prefix):
                text = text[len(prefix):]
                break

        # Strip articles
        for article in ["a ", "an ", "the "]:
            if text.startswith(article):
                text = text[len(article):]
                break

        text = text.strip().rstrip(".")

        return LABEL_MAP.get(text, text if text in LABEL_MAP.values() else None)

    # ── 3x3 Classification ──────────────────────────────────────────────────

    def preprocess_classifier(self, image: Image.Image) -> np.ndarray:
        """Preprocess image for per-class classifier: resize to 224x224, ImageNet normalize."""
        img = image.convert("RGB").resize((224, 224), Image.BILINEAR)
        arr = np.array(img, dtype=np.float32) / 255.0  # [H, W, 3]
        arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
        arr = np.transpose(arr, (2, 0, 1))  # -> [C, H, W]
        return np.expand_dims(arr, 0).astype(np.float32)  # -> [1, C, H, W]

    def classify_tile(self, image: Image.Image, label: str) -> float:
        """
        Classify a single tile image.
        Returns confidence score (0.0 - 1.0) that the tile contains the target object.
        """
        if label not in self.classifiers:
            logger.warning(f"No classifier for label: {label}")
            return 0.0

        session = self.classifiers[label]
        input_tensor = self.preprocess_classifier(image)

        # Get input name dynamically
        input_name = session.get_inputs()[0].name
        outputs = session.run(None, {input_name: input_tensor})
        logits = outputs[0][0]  # Shape: [2] for binary classifier

        # Softmax
        exp_logits = np.exp(logits - np.max(logits))
        probs = exp_logits / np.sum(exp_logits)

        return float(probs[1])  # Probability of target class (index 1)

    def classify_grid_3x3(
        self,
        tiles: List[Image.Image],
        label: str,
        is_hard: bool = False,
    ) -> List[int]:
        """
        Classify a 3x3 image grid.

        Args:
            tiles: List of 9 PIL Images (one per tile)
            label: Normalized object label
            is_hard: Whether this is a multi-step evaluation

        Returns:
            List of tile indices classified as positive matches (0-8)
        """
        threshold = 0.6 if is_hard else 0.7

        if len(tiles) != 9:
            logger.warning(f"Expected 9 tiles for 3x3, got {len(tiles)}")
            if len(tiles) < 9:
                return []

        # Classify each tile
        scores = []
        for i, tile in enumerate(tiles):
            score = self.classify_tile(tile, label)
            scores.append((i, score))
            logger.debug(f"Tile {i}: {label} = {score:.4f}")

        # Sort by confidence descending
        scores.sort(key=lambda x: x[1], reverse=True)

        # Select tiles above threshold
        selected = [i for i, s in scores if s > threshold]

        # Smart selection logic:
        # If 3 or 4 tiles pass threshold -> use them
        # Otherwise take top 3, but if gap between 3rd and 4th < 0.025, include 4th
        if len(selected) not in [3, 4]:
            if len(selected) < 3:
                selected = [scores[0][0], scores[1][0], scores[2][0]]
                if (
                    len(scores) >= 4
                    and scores[3][1] > 0
                    and (scores[2][1] - scores[3][1]) < 0.025
                ):
                    selected.append(scores[3][0])
            elif len(selected) > 4:
                # Take top 4
                selected = [scores[0][0], scores[1][0], scores[2][0], scores[3][0]]

        return sorted(selected)

    # ── 4x4 Segmentation ─────────────────────────────────────────────────────

    def preprocess_segmentation(self, image: Image.Image) -> np.ndarray:
        """Preprocess image for YOLOv5 segmentation: resize to 320x320, /255."""
        img = image.convert("RGB").resize((320, 320), Image.BILINEAR)
        arr = np.array(img, dtype=np.float32) / 255.0  # [H, W, 3]
        arr = np.transpose(arr, (2, 0, 1))  # -> [C, H, W]
        return np.expand_dims(arr, 0).astype(np.float32)  # -> [1, C, H, W]

    def classify_grid_4x4(self, background_image: Image.Image, label: str) -> List[int]:
        """
        Classify a 4x4 image grid using YOLOv5 instance segmentation.

        Model signatures (verified):
            yolov5-seg.ort:
                in:  images (1, 3, 320, 320)
                out: output0 (1, 6300, 52)  -- detections: 4 bbox + 1 obj + 15 cls + 32 mask
                     output1 (1, 32, 80, 80) -- mask prototypes

            nms-yolov5-det.ort:
                in:  detection (1, N, 52), config (3,) = [max_det, conf_thresh, iou_thresh]
                out: selected_idx (1, M) -- indices of surviving detections

            mask-yolov5-seg.ort:
                in:  detection (D,), mask (1, 32, 80, 80), config (9,)
                out: mask_filter (?, ?, 4) -- per-cell mask coverage

        Args:
            background_image: PIL Image of the full 4x4 grid (background image)
            label: Normalized object label

        Returns:
            List of cell indices classified as positive matches (0-15, row-major order)
        """
        if not self.seg_model:
            logger.error("Segmentation model not loaded, cannot process 4x4")
            return []

        input_tensor = self.preprocess_segmentation(background_image)

        # -- Step 1: Run YOLOv5 segmentation --
        seg_outputs = self.seg_model.run(
            None, {"images": input_tensor}
        )
        output0 = seg_outputs[0]  # (1, 6300, 52) -- raw detections
        output1 = seg_outputs[1]  # (1, 32, 80, 80) -- mask prototypes

        # -- Step 2: NMS to filter detections --
        if self.nms_model:
            nms_config = np.array([10, 0.25, 0.1], dtype=np.float32)
            try:
                nms_out = self.nms_model.run(
                    None,
                    {"detection": output0, "config": nms_config},
                )
                selected_idx = nms_out[0]  # (1, M) -- indices of kept detections
                if len(selected_idx.shape) == 2 and selected_idx.shape[1] > 0:
                    idx = selected_idx[0].astype(int)
                    kept = output0[0, idx, :]  # (M, 52)
                else:
                    kept = output0[0]
            except Exception as e:
                logger.warning(f"NMS failed, using raw detections: {e}")
                kept = output0[0]
        else:
            kept = output0[0]

        # -- Step 3: Filter by class and confidence --
        selected_cells = set()
        cell_size = 80.0  # 320 / 4
        label_idx = None
        for i, name in enumerate(YOLO_CLASSES):
            if name == label:
                label_idx = i
                break

        if label_idx is None:
            logger.error(f"Unknown label for 4x4: {label}")
            return []

        for det in kept:
            # Detection layout: [cx, cy, w, h, obj_conf, cls0..14, mask0..31]
            if len(det) < 52:
                continue

            obj_conf = float(det[4])
            if obj_conf < 0.25:
                continue

            class_scores = det[5:20]
            class_id = int(np.argmax(class_scores))
            if class_id != label_idx:
                continue

            confidence = float(class_scores[class_id]) * obj_conf
            logger.debug(f"Detection: {label} @ conf={confidence:.3f}")

            # Bounding box: [cx, cy, w, h] -> [x1, y1, x2, y2]
            cx, cy, w, h = float(det[0]), float(det[1]), float(det[2]), float(det[3])
            x1, y1 = cx - w / 2, cy - h / 2
            x2, y2 = cx + w / 2, cy + h / 2

            # -- Step 4: Mask processing (if mask model available) --
            if self.mask_model:
                try:
                    mask_coeffs = det[20:52].astype(np.float32)
                    mask_config = np.array([
                        80, 80, 320, 320,
                        x1, y1, x2, y2, 80.0,
                    ], dtype=np.float32)

                    mask_out = self.mask_model.run(
                        None,
                        {
                            "detection": mask_coeffs,
                            "mask": output1,
                            "config": mask_config,
                        },
                    )
                    mask_result = mask_out[0]
                    logger.debug(f"Mask output shape: {mask_result.shape}")

                    if mask_result.ndim >= 2 and mask_result.shape[0] == 4 and mask_result.shape[1] == 4:
                        for row in range(4):
                            for col in range(4):
                                coverage = float(mask_result[row, col])
                                if coverage > 0.10:
                                    selected_cells.add(row * 4 + col)
                        continue
                except Exception as e:
                    logger.debug(f"Mask model error: {e}")

            # -- Fallback: Compute mask manually from coefficients --
            try:
                mask_coeffs = det[20:52].astype(np.float32)
                mask_proto = output1[0]  # (32, 80, 80)

                mask = np.tensordot(mask_coeffs, mask_proto, axes=([0], [0]))
                mask = 1.0 / (1.0 + np.exp(-np.clip(mask, -500, 500)))  # Sigmoid

                bw = max(1, int(x2 - x1))
                bh = max(1, int(y2 - y1))
                mask_img = Image.fromarray((mask * 255).astype(np.uint8), mode="L")
                mask_img = mask_img.resize((bw, bh), Image.BILINEAR)
                mask_bin = np.array(mask_img) > 128

                for cell_row in range(4):
                    for cell_col in range(4):
                        cx1 = cell_col * cell_size
                        cy1 = cell_row * cell_size
                        cx2 = (cell_col + 1) * cell_size
                        cy2 = (cell_row + 1) * cell_size

                        mx1 = max(0, int(cx1 - x1))
                        my1 = max(0, int(cy1 - y1))
                        mx2 = min(mask_bin.shape[1], int(cx2 - x1))
                        my2 = min(mask_bin.shape[0], int(cy2 - y1))

                        if mx2 > mx1 and my2 > my1:
                            cell_mask = mask_bin[my1:my2, mx1:mx2]
                            total = (mx2 - mx1) * (my2 - my1)
                            coverage = np.sum(cell_mask) / total if total > 0 else 0
                            if coverage > 0.10:
                                selected_cells.add(cell_row * 4 + cell_col)
                continue

            except Exception:
                pass  # Fall through to bbox overlap

            # -- Final fallback: Bbox-based overlap (least accurate) --
            for cell_row in range(4):
                for cell_col in range(4):
                    cx1 = cell_col * cell_size
                    cy1 = cell_row * cell_size
                    cx2 = (cell_col + 1) * cell_size
                    cy2 = (cell_row + 1) * cell_size

                    ox1 = max(x1, cx1)
                    oy1 = max(y1, cy1)
                    ox2 = min(x2, cx2)
                    oy2 = min(y2, cy2)

                    if ox1 < ox2 and oy1 < oy2:
                        overlap = (ox2 - ox1) * (oy2 - oy1)
                        cell_area = cell_size * cell_size
                        if overlap / cell_area > 0.10:
                            selected_cells.add(cell_row * 4 + cell_col)

        result = sorted(selected_cells)
        logger.info(f"4x4 classification for '{label}': matched cells {result}")
        return result

    # ── Main Classification Entry ────────────────────────────────────────────

    def classify(
        self,
        task_description: str,
        tiles: List[Image.Image],
        background_image: Optional[Image.Image] = None,
        is_hard: bool = False,
    ) -> List[int]:
        """
        Main entry point for classifying an image grid.

        Args:
            task_description: Raw text describing the target object
            tiles: List of individual tile images (for 3x3)
            background_image: Full background image (for 4x4)
            is_hard: Whether this is a multi-step evaluation

        Returns:
            List of indices classified as positive matches
        """
        label = self.normalize_label(task_description)
        if not label:
            logger.error(f"Cannot parse label from: {task_description}")
            return []

        grid_size = len(tiles)
        logger.info(f"Classifying grid: label='{label}', size={grid_size}x{grid_size}, hard={is_hard}")

        if grid_size == 9 or grid_size == 16:
            if grid_size == 9 and label in self.classifiers:
                return self.classify_grid_3x3(tiles, label, is_hard)
            elif grid_size == 16 and self.seg_model:
                img = background_image if background_image else self._merge_tiles(tiles, 4)
                return self.classify_grid_4x4(img, label)
            elif grid_size == 9:
                if self.seg_model:
                    img = self._merge_tiles(tiles, 3)
                    return self.classify_grid_4x4(img, label)
                else:
                    logger.error(f"No model available for label: {label}")
                    return []
            elif grid_size == 16:
                logger.error("No segmentation model for 4x4")
                return []
        else:
            logger.warning(f"Unexpected grid size: {grid_size}")
            return []

    @staticmethod
    def _merge_tiles(tiles: List[Image.Image], grid_size: int) -> Image.Image:
        """Merge individual tile images into a single grid image."""
        cols = grid_size
        rows = grid_size
        tw, th = tiles[0].size

        merged = Image.new("RGB", (tw * cols, th * rows))
        for idx, tile in enumerate(tiles):
            r, c = divmod(idx, cols)
            merged.paste(tile.convert("RGB"), (c * tw, r * th))

        return merged