qwen_test_runner/vision/specialists.py

"""
specialists.py — the CPU assembly layer of the deterministic pipeline.

`Solids` holds one image's solidification primitives (detector boxes, seg masks, a
relative depth map, optional saliency, OCR, tags, class/style) — all in PIXEL space.
The `build_*` functions turn a `Solids` into each task's exact `tasks_vision` JSON, in
the task's declared coord space (via the existing `coords.from_canonical`), reusing the
`derive` engine for the INTEGRATE tasks.

The GPU half — loading the models and populating `Solids` — lives in the Colab runner
(`specialists_gpu`); this module is model-free and unit-tested on CPU.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Optional

import numpy as np

from . import derive
from .coords import BBox, CoordSpace, XYXY, _scale_for_space, from_canonical
from .tasks_vision import get_task


# ── coordinate helpers (pixel → task space) ──────────────────────────────────
def box_to_space(pix_xyxy, space: CoordSpace, size, ndigits: int = 2) -> list:
    """Pixel-abs [x1,y1,x2,y2] → the task's coord space."""
    return [round(v, ndigits) for v in from_canonical(BBox(*map(float, pix_xyxy)), space, size, XYXY)]


def poly_to_space(pix_flat, space: CoordSpace, size, ndigits: int = 2) -> list:
    """Flat pixel polygon [x1,y1,x2,y2,...] → the task's coord space, per vertex."""
    sx, sy = _scale_for_space(space, size)          # raw = pixel / scale
    out = []
    for i in range(0, len(pix_flat) - 1, 2):
        out.append(round(float(pix_flat[i]) / sx, ndigits))
        out.append(round(float(pix_flat[i + 1]) / sy, ndigits))
    return out


def quad_to_xyxy(quad) -> list:
    """A 4-point polygon (PaddleOCR) — [[x,y]*4] or flat [x,y,...] — → pixel xyxy."""
    a = np.asarray(quad, dtype=float).reshape(-1, 2)
    return [float(a[:, 0].min()), float(a[:, 1].min()), float(a[:, 0].max()), float(a[:, 1].max())]


# ── per-image primitives ─────────────────────────────────────────────────────
@dataclass
class Solids:
    size: tuple                       # (W, H) pixels
    boxes: list = field(default_factory=list)   # [{label, box:[x1,y1,x2,y2] px, score, mask?}]
    depth: Optional[np.ndarray] = None          # HxW relative (Depth-Anything: higher = nearer)
    saliency: Optional[np.ndarray] = None       # HxW float in [0,1]
    tags: list = field(default_factory=list)     # [{label, score}] (RAM++/SigLIP2)
    class_top: list = field(default_factory=list)  # [{label, score}] top-k classification
    ocr: Optional[dict] = None                   # {full_text, lines:[{text, box:[quad px], conf?}]}
    style: Optional[str] = None                  # SigLIP2 style label
    gray: Optional[np.ndarray] = None            # HxW grayscale (for symmetry)
    depth_higher_is_nearer: bool = True
    attr_boxes: list = field(default_factory=list)  # fusion tier: caption-phrase grounding
    #   [{phrase, matched_span, box:[x1,y1,x2,y2] px, score}] from ground_phrases()


# ── SOLID task builders (direct model outputs) ───────────────────────────────
# Emitted lists/strings are truncated to the registry caps (read from the spec,
# never hardcoded) — the generated pydantic models enforce max_items/max_length
# as hard constraints, so an uncapped emit would flunk schema validation on
# busy images even though the content is correct.
def build_bbox(s: Solids) -> dict:
    spec = get_task("bbox_grounding")
    cap = spec.fields["detections"].max_items
    boxes = sorted(s.boxes, key=lambda b: -float(b.get("score", 1.0)))[:cap]
    dets = [{"label": str(b["label"]),
             "box": box_to_space(b["box"], spec.coord_space, s.size),
             "score": round(float(b.get("score", 1.0)), 4)} for b in boxes]
    return {"detections": dets, "count": len(dets)}


def build_segmentation(s: Solids) -> dict:
    spec = get_task("segmentation")
    cap = spec.fields["masks"].max_items
    masked = [b for b in s.boxes if b.get("mask") is not None]
    masked.sort(key=lambda b: -float(np.asarray(b["mask"]).sum()))   # keep largest
    masks = []
    for b in masked:
        if len(masks) >= cap:
            break
        poly = derive.outline_polygon(b["mask"], b["label"])["outline"]
        if len(poly) >= 6:
            masks.append({"label": str(b["label"]),
                          "polygon": poly_to_space(poly, spec.coord_space, s.size)})
    return {"masks": masks}


def build_classification(s: Solids) -> dict:
    top = s.class_top or ([{"label": s.boxes[0]["label"], "score": 1.0}] if s.boxes else
                          [{"label": "unknown", "score": 0.0}])
    top = sorted(top, key=lambda t: -t["score"])[:5]
    return {"label": str(top[0]["label"]), "confidence": round(float(top[0]["score"]), 4),
            "top5": [{"label": str(t["label"]), "score": round(float(t["score"]), 4)} for t in top]}


def build_ocr(s: Solids) -> dict:
    spec = get_task("ocr_text")
    if not s.ocr:
        return {"full_text": "", "lines": []}
    lines_spec = spec.fields["lines"]
    text_cap = next((f.max_str_length for f in lines_spec.nested_fields
                     if f.name == "text"), 512)
    lines = []
    for ln in s.ocr.get("lines", [])[:lines_spec.max_items]:
        d = {"text": str(ln["text"])[:text_cap]}
        if ln.get("box") is not None:
            d["box"] = box_to_space(quad_to_xyxy(ln["box"]), spec.coord_space, s.size)
        lines.append(d)
    full_cap = spec.fields["full_text"].max_str_length
    return {"full_text": str(s.ocr.get("full_text", ""))[:full_cap], "lines": lines}


# ── INTEGRATE task builders (derive engine + coord conversion) ───────────────
def build_spatial(s: Solids) -> dict:
    return derive.spatial_relations(s.boxes, depth=s.depth,
                                    higher_is_nearer=s.depth_higher_is_nearer)


def build_depth_order(s: Solids) -> dict:
    if s.depth is None:
        return {"nearest": "", "farthest": "", "relative_depth": []}
    ents = [{"label": b["label"], "mask": b.get("mask"), "box": b["box"]} for b in s.boxes]
    return derive.depth_order(ents, s.depth, higher_is_nearer=s.depth_higher_is_nearer)


def build_subject(s: Solids) -> dict:
    space = get_task("subject_fixation").coord_space
    out = derive.subject_fixation(s.boxes, s.size, saliency=s.saliency)
    out["primary_subject"]["box"] = box_to_space(out["primary_subject"]["box"], space, s.size)
    return out


def build_outline(s: Solids) -> dict:
    space = get_task("outline_association").coord_space
    masked = [b for b in s.boxes if b.get("mask") is not None]
    if not masked:
        return {"outline": [], "label": ""}
    big = max(masked, key=lambda b: np.asarray(b["mask"]).sum())
    o = derive.outline_polygon(big["mask"], big["label"])
    o["outline"] = poly_to_space(o["outline"], space, s.size)
    return o


def build_style(s: Solids) -> dict:
    style = s.style if s.style in ("photo", "painting", "3d_render", "sketch", "anime", "other") else "other"
    layout = derive.layout_kind(s.boxes, s.size)
    symmetry = derive.symmetry_axis(s.gray) if s.gray is not None else "none"
    return {"style": style, "layout": layout, "symmetry": symmetry}


def build_datatype_diff(s: Solids) -> dict:
    return derive.detect_data_type(s.ocr.get("full_text", "") if s.ocr else "")


def build_datatype_util(s: Solids) -> tuple[dict, bool]:
    """Returns ({data_type, content}, parsed_ok). parsed_ok=False → caller may VLM-fallback."""
    return derive.parse_data_type(s.ocr.get("full_text", "") if s.ocr else "")


# task → builder (SOLID + INTEGRATE only; semantic_association + vqa are VLM)
DETERMINISTIC_BUILDERS = {
    "bbox_grounding": build_bbox,
    "segmentation": build_segmentation,
    "image_classification": build_classification,
    "ocr_text": build_ocr,
    "structural_spatial_awareness": build_spatial,
    "depth_analysis": build_depth_order,
    "subject_fixation": build_subject,
    "outline_association": build_outline,
    "style_structural_awareness": build_style,
    "data_type_differentiation": build_datatype_diff,
    # data_type_utilization handled specially (returns the parsed_ok flag)
}