benchmarks/modal/promptable_segmentation_shared.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import math
import random
from typing import Iterable

import numpy as np
from PIL import Image, ImageDraw, ImageFont


@dataclass(frozen=True)
class PromptableSceneObject:
    label: str
    prompt: str
    bbox: tuple[int, int, int, int]
    mask: np.ndarray
    color: tuple[int, int, int]


@dataclass(frozen=True)
class PromptableSceneFrame:
    image_path: str
    objects: list[PromptableSceneObject]
    target_index: int

    @property
    def target(self) -> PromptableSceneObject:
        return self.objects[self.target_index]


@dataclass(frozen=True)
class PromptableScene:
    scene_id: str
    prompt: str
    frames: list[PromptableSceneFrame]

    @property
    def reference_frame(self) -> PromptableSceneFrame:
        return self.frames[0]


def build_promptable_scene_suite(
    base_dir: Path,
    scene_count: int = 6,
    frame_count: int = 3,
    size: int = 512,
) -> list[PromptableScene]:
    """Generate synthetic shelf-like scenes for promptable-segmentation tests.

    The goal is not to mimic a specific retail image distribution. The goal
    is to isolate the promptability / mask quality / sweep stability behavior
    of the model family on a controlled object-selection task with known GT.
    """
    base_dir.mkdir(parents=True, exist_ok=True)
    rng = random.Random(20260614)
    palette = [
        ("red box", (214, 54, 54)),
        ("green pouch", (69, 156, 93)),
        ("blue carton", (70, 115, 214)),
        ("yellow bottle", (222, 180, 48)),
        ("purple jar", (156, 84, 198)),
        ("orange pack", (233, 138, 58)),
    ]
    scenes: list[PromptableScene] = []

    for scene_idx in range(scene_count):
        scene_dir = base_dir / f"scene_{scene_idx:02d}"
        scene_dir.mkdir(parents=True, exist_ok=True)
        prompt_label, prompt_color = palette[scene_idx % len(palette)]
        prompt = prompt_label

        frames: list[PromptableSceneFrame] = []
        base_objects = _make_scene_objects(
            prompt_label=prompt_label,
            prompt_color=prompt_color,
            rng=rng,
            size=size,
        )

        for frame_idx in range(frame_count):
            frame_path = scene_dir / f"frame_{frame_idx:02d}.png"
            objects = _render_scene_frame(
                frame_path=frame_path,
                size=size,
                objects=base_objects,
                frame_idx=frame_idx,
                scene_idx=scene_idx,
            )
            target_index = next(
                idx for idx, obj in enumerate(objects) if obj.label == prompt_label
            )
            frames.append(
                PromptableSceneFrame(
                    image_path=str(frame_path),
                    objects=objects,
                    target_index=target_index,
                )
            )
        scenes.append(PromptableScene(scene_id=f"scene_{scene_idx:02d}", prompt=prompt, frames=frames))

    return scenes


def mask_iou(pred: np.ndarray, gt: np.ndarray) -> float:
    pred_bin = _as_bool_mask(pred)
    gt_bin = _as_bool_mask(gt)
    union = np.logical_or(pred_bin, gt_bin).sum()
    if union == 0:
        return 0.0
    inter = np.logical_and(pred_bin, gt_bin).sum()
    return float(inter / union)


def pixel_accuracy(pred: np.ndarray, gt: np.ndarray) -> float:
    pred_bin = _as_bool_mask(pred)
    gt_bin = _as_bool_mask(gt)
    return float((pred_bin == gt_bin).mean())


def centroid_px(mask: np.ndarray) -> tuple[float, float] | None:
    ys, xs = np.where(_as_bool_mask(mask))
    if xs.size == 0 or ys.size == 0:
        return None
    return (float(xs.mean()), float(ys.mean()))


def bbox_iou(box_a: Iterable[float], box_b: Iterable[float]) -> float:
    ax1, ay1, ax2, ay2 = [float(v) for v in box_a]
    bx1, by1, bx2, by2 = [float(v) for v in box_b]
    ix1 = max(ax1, bx1)
    iy1 = max(ay1, by1)
    ix2 = min(ax2, bx2)
    iy2 = min(ay2, by2)
    inter_w = max(0.0, ix2 - ix1)
    inter_h = max(0.0, iy2 - iy1)
    inter = inter_w * inter_h
    if inter <= 0:
        return 0.0
    area_a = max(0.0, ax2 - ax1) * max(0.0, ay2 - ay1)
    area_b = max(0.0, bx2 - bx1) * max(0.0, by2 - by1)
    denom = area_a + area_b - inter
    return float(inter / denom) if denom > 0 else 0.0


def mask_from_bbox(shape: tuple[int, int], bbox: Iterable[float]) -> np.ndarray:
    height, width = shape
    x1, y1, x2, y2 = [int(round(v)) for v in bbox]
    arr = np.zeros((height, width), dtype=bool)
    arr[max(0, y1) : min(height, y2), max(0, x1) : min(width, x2)] = True
    return arr


def best_mask_from_prediction(prediction: object, image_size: tuple[int, int]) -> np.ndarray | None:
    """Extract the first usable binary mask from a candidate prediction."""
    height, width = image_size

    # Ultralytics-style Results objects
    masks = getattr(prediction, "masks", None)
    if masks is not None:
        data = getattr(masks, "data", None)
        if data is not None and len(data):
            first = data[0]
            if hasattr(first, "detach"):
                first = first.detach()
            if hasattr(first, "cpu"):
                first = first.cpu()
            if hasattr(first, "numpy"):
                arr = first.numpy()
            else:
                arr = np.asarray(first)
            if arr.ndim == 3:
                arr = arr[0]
            return _resize_bool(arr, (height, width))

        xy = getattr(masks, "xy", None)
        if xy:
            poly = xy[0]
            return _polygon_to_mask(poly, (height, width))

    # RF-DETR / detector-style results may expose boxes + masks.
    boxes = getattr(prediction, "boxes", None)
    if boxes is not None:
        xyxy = getattr(boxes, "xyxy", None)
        if xyxy is not None and len(xyxy):
            first = xyxy[0]
            if hasattr(first, "cpu"):
                first = first.cpu()
            if hasattr(first, "numpy"):
                bbox = first.numpy()
            else:
                bbox = np.asarray(first)
            return mask_from_bbox((height, width), bbox)

    return None


def promptability_score(
    pred_mask: np.ndarray | None,
    gt_mask: np.ndarray,
    gt_box: Iterable[float],
) -> dict[str, float]:
    if pred_mask is None:
        return {
            "mask_iou": 0.0,
            "pixel_acc": 0.0,
            "bbox_iou": 0.0,
            "centroid_error_px": float("inf"),
        }
    pred_box = _mask_bbox(pred_mask)
    gt_centroid = centroid_px(gt_mask)
    pred_centroid = centroid_px(pred_mask)
    centroid_error = (
        float(math.dist(gt_centroid, pred_centroid))
        if gt_centroid is not None and pred_centroid is not None
        else float("inf")
    )
    return {
        "mask_iou": mask_iou(pred_mask, gt_mask),
        "pixel_acc": pixel_accuracy(pred_mask, gt_mask),
        "bbox_iou": bbox_iou(pred_box, gt_box),
        "centroid_error_px": centroid_error,
    }


def _make_scene_objects(
    prompt_label: str,
    prompt_color: tuple[int, int, int],
    rng: random.Random,
    size: int,
) -> list[PromptableSceneObject]:
    labels = [
        (prompt_label, prompt_color),
        ("gray tin", (136, 136, 136)),
        ("teal pouch", (58, 153, 148)),
    ]
    objects: list[PromptableSceneObject] = []
    occupied: list[tuple[int, int, int, int]] = []
    for label, color in labels:
        bbox = _place_box(rng, size=size, occupied=occupied)
        occupied.append(bbox)
        objects.append(
            PromptableSceneObject(
                label=label,
                prompt=label,
                bbox=bbox,
                mask=_rect_mask((size, size), bbox),
                color=color,
            )
        )
    return objects


def _render_scene_frame(
    frame_path: Path,
    size: int,
    objects: list[PromptableSceneObject],
    frame_idx: int,
    scene_idx: int,
) -> list[PromptableSceneObject]:
    image = Image.new("RGB", (size, size), (250, 249, 246))
    draw = ImageDraw.Draw(image)
    _draw_shelf_grid(draw, size=size)
    shifted: list[PromptableSceneObject] = []

    for obj_idx, obj in enumerate(objects):
        dx = (frame_idx * 4) + (scene_idx % 3) * 2
        dy = ((frame_idx + obj_idx) % 2) * 3
        jittered = _shift_box(obj.bbox, dx=dx, dy=dy, size=size)
        mask = _rect_mask((size, size), jittered)
        shifted.append(
            PromptableSceneObject(
                label=obj.label,
                prompt=obj.prompt,
                bbox=jittered,
                mask=mask,
                color=obj.color,
            )
        )
        _draw_object(draw, jittered, obj.label, obj.color)

    image.save(frame_path)
    return shifted


def _draw_shelf_grid(draw: ImageDraw.ImageDraw, size: int) -> None:
    shelf_y = int(size * 0.21)
    for offset in [shelf_y, int(size * 0.5), int(size * 0.79)]:
        draw.rectangle([24, offset, size - 24, offset + 10], fill=(206, 199, 189))


def _draw_object(
    draw: ImageDraw.ImageDraw,
    bbox: tuple[int, int, int, int],
    label: str,
    color: tuple[int, int, int],
) -> None:
    x1, y1, x2, y2 = bbox
    draw.rounded_rectangle([x1, y1, x2, y2], radius=16, fill=color, outline=(40, 40, 40), width=3)
    try:
        font = ImageFont.truetype("/System/Library/Fonts/Supplemental/Arial Bold.ttf", 18)
    except Exception:
        font = ImageFont.load_default()
    text_bbox = draw.textbbox((0, 0), label, font=font)
    tw = text_bbox[2] - text_bbox[0]
    tx = x1 + max(8, (x2 - x1 - tw) // 2)
    ty = max(6, y1 + 8)
    draw.rounded_rectangle([x1 + 6, y1 + 6, min(x2 - 6, x1 + 6 + tw + 12), y1 + 32], radius=8, fill=(255, 255, 255))
    draw.text((tx, ty), label, fill=(20, 20, 20), font=font)


def _place_box(
    rng: random.Random,
    size: int,
    occupied: list[tuple[int, int, int, int]],
) -> tuple[int, int, int, int]:
    for _ in range(200):
        w = rng.randint(int(size * 0.18), int(size * 0.28))
        h = rng.randint(int(size * 0.20), int(size * 0.32))
        x1 = rng.randint(28, size - w - 28)
        y1 = rng.randint(28, size - h - 28)
        candidate = (x1, y1, x1 + w, y1 + h)
        if all(_iou(candidate, other) < 0.08 for other in occupied):
            return candidate
    return (60, 80, 60 + int(size * 0.22), 80 + int(size * 0.24))


def _shift_box(
    bbox: tuple[int, int, int, int],
    dx: int,
    dy: int,
    size: int,
) -> tuple[int, int, int, int]:
    x1, y1, x2, y2 = bbox
    width = x2 - x1
    height = y2 - y1
    nx1 = min(max(20, x1 + dx), size - width - 20)
    ny1 = min(max(20, y1 + dy), size - height - 20)
    return (nx1, ny1, nx1 + width, ny1 + height)


def _rect_mask(shape: tuple[int, int], bbox: tuple[int, int, int, int]) -> np.ndarray:
    height, width = shape
    mask = np.zeros((height, width), dtype=bool)
    x1, y1, x2, y2 = bbox
    mask[max(0, y1) : min(height, y2), max(0, x1) : min(width, x2)] = True
    return mask


def _resize_bool(mask: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    image = Image.fromarray((_as_bool_mask(mask).astype(np.uint8) * 255))
    resized = image.resize((size[1], size[0]), Image.NEAREST)
    return np.array(resized) > 127


def _polygon_to_mask(poly: np.ndarray, size: tuple[int, int]) -> np.ndarray:
    image = Image.new("L", (size[1], size[0]), 0)
    draw = ImageDraw.Draw(image)
    if hasattr(poly, "tolist"):
        poly = poly.tolist()
    if poly and isinstance(poly[0], (list, tuple)):
        points = [(float(x), float(y)) for x, y in poly]
        draw.polygon(points, outline=1, fill=1)
    return np.array(image, dtype=bool)


def _mask_bbox(mask: np.ndarray) -> tuple[float, float, float, float]:
    ys, xs = np.where(_as_bool_mask(mask))
    if xs.size == 0 or ys.size == 0:
        return (0.0, 0.0, 0.0, 0.0)
    return (float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max()))


def _as_bool_mask(mask: np.ndarray) -> np.ndarray:
    arr = np.asarray(mask)
    if arr.dtype == bool:
        return arr
    if arr.ndim > 2:
        arr = arr.squeeze()
    return arr > 0.5


def _iou(a: tuple[int, int, int, int], b: tuple[int, int, int, int]) -> float:
    ax1, ay1, ax2, ay2 = a
    bx1, by1, bx2, by2 = b
    ix1 = max(ax1, bx1)
    iy1 = max(ay1, by1)
    ix2 = min(ax2, bx2)
    iy2 = min(ay2, by2)
    inter = max(0, ix2 - ix1) * max(0, iy2 - iy1)
    if inter <= 0:
        return 0.0
    area_a = max(0, ax2 - ax1) * max(0, ay2 - ay1)
    area_b = max(0, bx2 - bx1) * max(0, by2 - by1)
    denom = area_a + area_b - inter
    return inter / denom if denom else 0.0