qwen_test_runner/vision/coords.py

"""
coords.py — Coordinate normalization (the correctness centerpiece).

Different VLMs report boxes in different spaces:
  * Qwen3-VL emits RELATIVE coordinates in 0..1000 (integers).
  * Qwen3.5 / many checkpoints emit 0..1 floats.
  * COCO / LVIS ground truth is ABSOLUTE pixels.

If predictions and ground truth are compared in mismatched spaces, every IoU is
silently wrong and detection/3D/segmentation metrics collapse. To prevent that,
the canonical internal form is ALWAYS pixel-absolute xyxy. GT is converted to
canonical at load time; predictions are converted at score time. No metric ever
sees a non-canonical coordinate (enforced by tests).
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Sequence


class CoordSpace(str, Enum):
    """The space a set of raw coordinates lives in."""
    PIXEL_ABS = "pixel_abs"      # absolute pixels, image-sized
    NORM_0_1 = "norm_0_1"        # 0..1 floats
    NORM_0_1000 = "norm_0_1000"  # 0..1000 ints (Qwen3-VL native)


# Box coordinate layouts a model might emit / GT might store.
XYXY = "xyxy"  # [x1, y1, x2, y2]
XYWH = "xywh"  # [x, y, w, h]  (COCO GT layout)


def xywh_to_xyxy(box: Sequence[float]) -> list[float]:
    x, y, w, h = box
    return [x, y, x + w, y + h]


def xyxy_to_xywh(box: Sequence[float]) -> list[float]:
    x1, y1, x2, y2 = box
    return [x1, y1, x2 - x1, y2 - y1]


@dataclass(frozen=True)
class BBox:
    """A bounding box in the canonical form: pixel-absolute xyxy."""
    x1: float
    y1: float
    x2: float
    y2: float

    def area(self) -> float:
        return max(0.0, self.x2 - self.x1) * max(0.0, self.y2 - self.y1)

    def clip(self, size: tuple[int, int]) -> "BBox":
        """Clip to image bounds. size = (W, H)."""
        w, h = size
        return BBox(
            x1=min(max(self.x1, 0.0), w),
            y1=min(max(self.y1, 0.0), h),
            x2=min(max(self.x2, 0.0), w),
            y2=min(max(self.y2, 0.0), h),
        )

    def iou(self, other: "BBox") -> float:
        ix1 = max(self.x1, other.x1)
        iy1 = max(self.y1, other.y1)
        ix2 = min(self.x2, other.x2)
        iy2 = min(self.y2, other.y2)
        iw = max(0.0, ix2 - ix1)
        ih = max(0.0, iy2 - iy1)
        inter = iw * ih
        if inter <= 0.0:
            return 0.0
        union = self.area() + other.area() - inter
        return inter / union if union > 0 else 0.0

    def as_list(self) -> list[float]:
        return [self.x1, self.y1, self.x2, self.y2]


def _scale_for_space(space: CoordSpace, size: tuple[int, int]) -> tuple[float, float]:
    """Return (x_scale, y_scale) that maps a raw coord in `space` to pixels."""
    w, h = size
    if space == CoordSpace.PIXEL_ABS:
        return 1.0, 1.0
    if space == CoordSpace.NORM_0_1:
        return float(w), float(h)
    if space == CoordSpace.NORM_0_1000:
        return w / 1000.0, h / 1000.0
    raise ValueError(f"unknown coord space: {space!r}")


def to_canonical(
    raw: Sequence[float],
    space: CoordSpace,
    size: tuple[int, int],
    fmt: str = XYXY,
) -> BBox:
    """Convert a raw 4-tuple in `space`/`fmt` to a canonical pixel-abs xyxy BBox.

    size = (W, H) in pixels. `fmt` is XYXY or XYWH.
    """
    if len(raw) != 4:
        raise ValueError(f"bbox must have 4 values, got {len(raw)}: {raw!r}")
    coords = list(map(float, raw))
    if fmt == XYWH:
        coords = xywh_to_xyxy(coords)
    elif fmt != XYXY:
        raise ValueError(f"unknown bbox fmt: {fmt!r}")
    sx, sy = _scale_for_space(space, size)
    return BBox(coords[0] * sx, coords[1] * sy, coords[2] * sx, coords[3] * sy).clip(size)


def from_canonical(box: BBox, space: CoordSpace, size: tuple[int, int], fmt: str = XYXY) -> list[float]:
    """Inverse of to_canonical: canonical BBox → raw coords in `space`/`fmt`."""
    sx, sy = _scale_for_space(space, size)
    xyxy = [box.x1 / sx, box.y1 / sy, box.x2 / sx, box.y2 / sy]
    return xyxy_to_xywh(xyxy) if fmt == XYWH else xyxy


def detect_space(raw_values: Sequence[float], size: tuple[int, int]) -> CoordSpace:
    """Defensive fallback for models that ignore the requested space.

    Used ONLY when a model's output space can't be trusted; the caller logs
    `coord_space_inferred=True` (itself a robustness signal). Heuristic:
      * all values <= 1.0           → NORM_0_1
      * all values <= 1000 and the image is larger than 1000 px on a side → NORM_0_1000
      * otherwise                   → PIXEL_ABS
    """
    vals = [abs(float(v)) for v in raw_values if v is not None]
    if not vals:
        return CoordSpace.PIXEL_ABS
    mx = max(vals)
    w, h = size
    if mx <= 1.0:
        return CoordSpace.NORM_0_1
    if mx <= 1000.0 and max(w, h) > 1000:
        return CoordSpace.NORM_0_1000
    if mx <= 1000.0 and max(w, h) <= 1000:
        # ambiguous: 0..1000 ints vs small-image pixels. Prefer NORM_0_1000 only
        # if values clearly exceed the image dimensions.
        if mx > max(w, h):
            return CoordSpace.NORM_0_1000
        return CoordSpace.PIXEL_ABS
    return CoordSpace.PIXEL_ABS


def prompt_hint_for(space: CoordSpace) -> str:
    """A sentence appended to the system prompt telling the model which space to use."""
    if space == CoordSpace.PIXEL_ABS:
        return "Report bounding boxes as [x1, y1, x2, y2] in absolute pixel coordinates."
    if space == CoordSpace.NORM_0_1:
        return "Report bounding boxes as [x1, y1, x2, y2] normalized to 0..1 of the image dimensions."
    if space == CoordSpace.NORM_0_1000:
        return (
            "Report bounding boxes as [x1, y1, x2, y2] integers in 0..1000, "
            "relative to the image width and height."
        )
    raise ValueError(f"unknown coord space: {space!r}")