qwen_test_runner/vision/derive.py

"""
derive.py — the deterministic "semantic engine": derive INTEGRATE tasks from the
solidification primitives (detector boxes, segmentation masks, a relative depth map,
optional saliency), with NO model. Pure numpy + stdlib; OpenCV is used lazily only for
the outline contour.

Every function returns the exact JSON shape of its `tasks_vision` task, so the output
validates against the task's registry Pydantic model and scores through the existing
`score_vision_sample`. Boxes/polygons are in whatever coordinate space the caller passes
in (pixels for real specialists) — coord-space normalization is done by the adapter layer,
not here.

Primitive conventions (documented, unit-tested):
  box   = [x1, y1, x2, y2]  (x right, y DOWN)
  mask  = HxW bool ndarray
  depth = HxW float ndarray, relative/ordinal. Depth-Anything convention: HIGHER = NEARER
          (disparity-like). Pass higher_is_nearer=False for metric depth (smaller = nearer).
"""

from __future__ import annotations

import json
import re
from typing import Optional, Sequence

import numpy as np

# tasks_vision closed vocabularies (kept in sync with the registry)
_DATATYPE_VALUES = ("json", "yaml", "markdown", "csv", "toml", "xml", "code", "plaintext")


# ── geometry helpers ─────────────────────────────────────────────────────────
def _centroid(b):
    return (0.5 * (b[0] + b[2]), 0.5 * (b[1] + b[3]))


def _area(b):
    return max(0.0, b[2] - b[0]) * max(0.0, b[3] - b[1])


def _iou(a, b):
    ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
    ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
    iw, ih = max(0.0, ix2 - ix1), max(0.0, iy2 - iy1)
    inter = iw * ih
    u = _area(a) + _area(b) - inter
    return inter / u if u > 0 else 0.0


def _contains(outer, inner, frac=0.85):
    """True if `inner` is (mostly) inside `outer` — inter/area(inner) >= frac."""
    ix1, iy1 = max(outer[0], inner[0]), max(outer[1], inner[1])
    ix2, iy2 = min(outer[2], inner[2]), min(outer[3], inner[3])
    inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
    ai = _area(inner)
    return ai > 0 and inter / ai >= frac


def _uniq_labels(labels):
    """Disambiguate duplicate label strings: person, person -> person_1, person_2."""
    seen, counts = {}, {}
    for l in labels:
        counts[l] = counts.get(l, 0) + 1
    out, running = [], {}
    for l in labels:
        if counts[l] == 1:
            out.append(l)
        else:
            running[l] = running.get(l, 0) + 1
            out.append(f"{l}_{running[l]}")
    return out


# ── #8 structural_spatial_awareness ──────────────────────────────────────────
def spatial_relations(boxes: Sequence[dict], depth: Optional[np.ndarray] = None,
                      higher_is_nearer: bool = True, max_items: int = 12) -> dict:
    """boxes: [{label, box:[x1,y1,x2,y2], score?}]. Emits projective relations
    (left_of/right_of/above/below), containment (inside), and — if a depth map is given —
    in_front_of/behind. Returns {"relations":[{subject,predicate,object}]}."""
    labs = _uniq_labels([str(b["label"]) for b in boxes])
    bxs = [b["box"] for b in boxes]
    n = len(bxs)
    # per-box depth (median over the box region) if a map is provided
    box_depth = None
    if depth is not None and n:
        H, W = depth.shape[:2]
        box_depth = []
        for b in bxs:
            x1, y1, x2, y2 = (int(round(b[0])), int(round(b[1])), int(round(b[2])), int(round(b[3])))
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(W, max(x1 + 1, x2)), min(H, max(y1 + 1, y2))
            patch = depth[y1:y2, x1:x2]
            box_depth.append(float(np.median(patch)) if patch.size else 0.0)

    rels, seen = [], set()   # seen holds (subject, predicate, object) triples
    # order pairs by centroid distance so the nearest, most meaningful pairs win the budget
    cents = [_centroid(b) for b in bxs]
    pairs = [(i, j) for i in range(n) for j in range(n) if i != j]
    pairs.sort(key=lambda ij: (cents[ij[0]][0] - cents[ij[1]][0]) ** 2
               + (cents[ij[0]][1] - cents[ij[1]][1]) ** 2)

    def _emit(si, pred, oi):
        t = (labs[si], pred, labs[oi])
        if t not in seen:
            seen.add(t)
            rels.append({"subject": labs[si], "predicate": pred, "object": labs[oi]})

    for i, j in pairs:
        if len(rels) >= max_items:
            break
        a, ca, cb = bxs[i], cents[i], cents[j]
        if _contains(bxs[j], a):                          # containment first (most specific)
            _emit(i, "inside", j)
            continue
        dx, dy = cb[0] - ca[0], cb[1] - ca[1]
        if abs(dx) >= abs(dy):
            _emit(i, "left_of" if dx > 0 else "right_of", j)   # a left_of b when a.x < b.x
        else:
            _emit(i, "above" if dy > 0 else "below", j)        # y DOWN: a above b when a.y < b.y
    # depth relations — a pair can carry BOTH a projective and a depth relation
    if box_depth is not None:
        rng = (max(box_depth) - min(box_depth)) or 1.0
        for i, j in pairs:
            if len(rels) >= max_items:
                break
            d = (box_depth[i] - box_depth[j]) / rng
            if abs(d) < 0.15:
                continue
            a_nearer = (d > 0) if higher_is_nearer else (d < 0)
            _emit(i, "in_front_of" if a_nearer else "behind", j)
    return {"relations": rels}


# ── #7 depth_analysis (ordering) ─────────────────────────────────────────────
def depth_scalars(entities: Sequence[dict], depth: np.ndarray,
                  higher_is_nearer: bool = True) -> list[float]:
    """Continuous per-entity NEARNESS in [0,1] (bigger = nearer): median relative depth
    over each entity's mask (preferred) or box, min-max normalized across the entities.
    This is the scalar core of `depth_order`, exposed so the fusion tier can keep the
    continuous signal instead of only the categorical nearer/farther/same."""
    if not entities:
        return []
    H, W = depth.shape[:2]
    vals = []
    for e in entities:
        if e.get("mask") is not None:
            m = np.asarray(e["mask"], dtype=bool)
            v = float(np.median(depth[m])) if m.any() else 0.0
        else:
            b = e["box"]
            x1, y1 = max(0, int(b[0])), max(0, int(b[1]))
            x2, y2 = min(W, int(b[2])), min(H, int(b[3]))
            patch = depth[y1:max(y1 + 1, y2), x1:max(x1 + 1, x2)]
            v = float(np.median(patch)) if patch.size else 0.0
        vals.append(v)
    # normalize to [0,1] for a stable "same" tolerance
    lo, hi = min(vals), max(vals)
    rng = (hi - lo) or 1.0
    norm = [(v - lo) / rng for v in vals]
    # nearness score: bigger = nearer
    return norm if higher_is_nearer else [1 - x for x in norm]


def depth_order(entities: Sequence[dict], depth: np.ndarray,
                higher_is_nearer: bool = True, same_tol: float = 0.08,
                max_items: int = 12) -> dict:
    """entities: [{label, mask:HxW bool}] (preferred) or [{label, box}]. Samples the
    RELATIVE depth over each entity (mask median, foreground-robust), orders them, and
    emits {"nearest","farthest","relative_depth":[{a,b,a_is}]}."""
    if not entities:
        return {"nearest": "", "farthest": "", "relative_depth": []}
    labs = _uniq_labels([str(e["label"]) for e in entities])
    near = depth_scalars(entities, depth, higher_is_nearer)
    order = sorted(range(len(labs)), key=lambda i: -near[i])
    out = {"nearest": labs[order[0]], "farthest": labs[order[-1]], "relative_depth": []}
    for i in range(len(labs)):
        for j in range(i + 1, len(labs)):
            if len(out["relative_depth"]) >= max_items:
                return out
            d = near[i] - near[j]
            a_is = "same" if abs(d) < same_tol else ("nearer" if d > 0 else "farther")
            out["relative_depth"].append({"a": labs[i], "b": labs[j], "a_is": a_is})
    return out


# ── #9 subject_fixation ──────────────────────────────────────────────────────
def subject_scores(boxes: Sequence[dict], image_size,
                   saliency: Optional[np.ndarray] = None) -> list[float]:
    """Per-box subject score (saliency-PRIMARY, area×centrality tie-break) for EVERY
    box — the scoring core of `subject_fixation`, exposed so the fusion tier can keep
    the full ranking instead of only the winner."""
    W, H = image_size
    cx, cy = W / 2.0, H / 2.0
    diag = (W ** 2 + H ** 2) ** 0.5 or 1.0

    def score(b):
        bx = b["box"]
        area = _area(bx) / (W * H + 1e-9)
        ctr = _centroid(bx)
        centrality = 1.0 - (((ctr[0] - cx) ** 2 + (ctr[1] - cy) ** 2) ** 0.5) / diag
        geo = area * (0.5 + 0.5 * centrality)
        if saliency is not None:
            x1, y1 = max(0, int(bx[0])), max(0, int(bx[1]))
            x2, y2 = min(int(saliency.shape[1]), int(bx[2])), min(int(saliency.shape[0]), int(bx[3]))
            patch = saliency[y1:max(y1 + 1, y2), x1:max(x1 + 1, x2)]
            sal = float(patch.mean()) if patch.size else 0.0
            return sal + 0.01 * geo          # saliency primary, geometry breaks ties
        return geo

    return [score(b) for b in boxes]


def subject_fixation(boxes: Sequence[dict], image_size, saliency: Optional[np.ndarray] = None) -> dict:
    """Saliency-PRIMARY (mean saliency inside each box), area×centrality tie-break.
    image_size = (W, H). Falls back to the largest box, then whole-image. Returns
    {"primary_subject":{"label","box"}}."""
    W, H = image_size
    if not boxes:
        return {"primary_subject": {"label": "scene", "box": [0.0, 0.0, float(W), float(H)]}}
    scores = subject_scores(boxes, image_size, saliency)
    best = boxes[max(range(len(boxes)), key=scores.__getitem__)]
    return {"primary_subject": {"label": str(best["label"]),
                                "box": [float(v) for v in best["box"]]}}


# ── #10 outline_association (mask → contour polygon) ─────────────────────────
def outline_polygon(mask: np.ndarray, label: str, max_points: int = 128) -> dict:
    """SAM2 mask → closed outline polygon, flat [x1,y1,x2,y2,...]. Pure-numpy row-scan:
    trace the left boundary top→bottom, then the right boundary bottom→top (a closed loop).
    No OpenCV dependency; the dense boundary is IoU-accurate (subsampled to max_points)."""
    m = np.asarray(mask) > 0
    if m.ndim != 2 or not m.any():
        return {"outline": [], "label": str(label)}
    ys = np.where(m.any(axis=1))[0]
    left, right = [], []
    for y in ys:
        xs = np.where(m[y])[0]
        left.append((float(xs[0]), float(y)))
        right.append((float(xs[-1]), float(y)))
    pts = left + right[::-1]                            # closed loop
    if len(pts) > max_points:
        idx = np.linspace(0, len(pts) - 1, max_points).round().astype(int)
        pts = [pts[i] for i in idx]
    flat = [v for xy in pts for v in xy]               # (x, y) interleaved
    return {"outline": flat, "label": str(label)}


# ── #6 style: symmetry + layout (the deterministic halves) ───────────────────
def symmetry_scores(gray: np.ndarray) -> dict:
    """Continuous L/R and T/B mirror correlations in [-1,1] — the scalar core of
    `symmetry_axis`, exposed so the fusion tier keeps the magnitudes the categorical
    label throws away. Returns {"lr": float, "tb": float}."""
    g = np.asarray(gray, dtype=np.float64)
    if g.ndim == 3:
        g = g.mean(axis=2)
    g = g - g.mean()

    def corr(a, b):
        a, b = a.ravel(), b.ravel()
        da, db = np.linalg.norm(a), np.linalg.norm(b)
        return float(a @ b / (da * db)) if da > 0 and db > 0 else 0.0

    return {"lr": corr(g, g[:, ::-1]),     # left-right mirror  -> vertical-axis symmetry
            "tb": corr(g, g[::-1, :])}     # top-bottom mirror  -> horizontal-axis symmetry


def symmetry_axis(gray: np.ndarray, thresh: float = 0.80) -> str:
    """Normalized-correlation of the image vs its L/R and T/B flips. Returns one of
    horizontal/vertical/radial/none. 'vertical' = mirror across a vertical axis (L==R)."""
    s = symmetry_scores(gray)
    lr, tb = s["lr"], s["tb"]
    v, h = lr >= thresh, tb >= thresh
    if v and h:
        return "radial"
    if v:
        return "vertical"
    if h:
        return "horizontal"
    return "none"


def layout_kind(boxes: Sequence[dict], image_size) -> str:
    """From the box constellation: centered / rule_of_thirds / symmetric / scattered /
    unknown."""
    W, H = image_size
    if not boxes:
        return "unknown"
    cents = np.array([_centroid(b["box"]) for b in boxes], dtype=float)
    areas = np.array([_area(b["box"]) for b in boxes], dtype=float)
    if len(boxes) == 1 or areas.max() > 0.5 * (W * H):
        cx, cy = cents[int(areas.argmax())]
        if abs(cx - W / 2) < 0.15 * W and abs(cy - H / 2) < 0.15 * H:
            return "centered"
    # left-right centroid symmetry about the vertical axis
    xs = cents[:, 0] / W
    if len(xs) >= 2 and abs(np.mean(xs) - 0.5) < 0.08 and np.std(xs) > 0.15:
        return "symmetric"
    # proximity to rule-of-thirds lines
    thirds = np.array([1 / 3, 2 / 3])
    nx = np.min(np.abs((cents[:, 0] / W)[:, None] - thirds[None, :]), axis=1)
    ny = np.min(np.abs((cents[:, 1] / H)[:, None] - thirds[None, :]), axis=1)
    if np.mean((nx < 0.08) | (ny < 0.08)) > 0.5:
        return "rule_of_thirds"
    return "scattered"


# ── #12/#13 data-type recognition + re-serialization ─────────────────────────
def detect_data_type(text: str) -> dict:
    """OCR text -> {"data_type","confidence"}. Deterministic regex/heuristic with a
    precedence order and a confidence proxy."""
    t = (text or "").strip()
    if not t:
        return {"data_type": "plaintext", "confidence": 0.2}
    scores = {k: 0.0 for k in _DATATYPE_VALUES}
    if re.search(r"^\s*[{\[]", t) and re.search(r"[}\]]\s*$", t) and '"' in t:
        scores["json"] += 0.9
    if re.search(r"^\s*<\?xml|</[a-zA-Z]", t):
        scores["xml"] += 0.9
    if re.search(r"^\s*---\s*$", t, re.M) or re.search(r"^\s*[\w-]+:\s+\S", t, re.M):
        scores["yaml"] += 0.6
    if re.search(r"^\s*#{1,6}\s|\*\*|\[.+\]\(.+\)|^\s*[-*]\s", t, re.M):
        scores["markdown"] += 0.6
    if re.search(r"^\s*\[[\w.\-]+\]\s*$", t, re.M) or re.search(r'^\s*[\w.-]+\s*=\s*("|\d|\[)', t, re.M):
        scores["toml"] += 0.6
    if "\n" in t and all("," in ln for ln in t.splitlines()[:3] if ln.strip()):
        scores["csv"] += 0.5
    if re.search(r"\b(def|function|class|import|return|const|var|let)\b", t):
        scores["code"] += 0.4
    best = max(scores, key=scores.get)
    conf = scores[best]
    if conf <= 0.0:
        return {"data_type": "plaintext", "confidence": 0.3}
    return {"data_type": best, "confidence": round(min(0.99, conf), 2)}


def _repair_json(t: str) -> str:
    t = t.strip().strip("`")
    t = re.sub(r"[“”]", '"', t)          # curly double quotes
    t = re.sub(r"[‘’]", "'", t)          # curly single quotes
    t = re.sub(r",\s*([}\]])", r"\1", t)           # trailing commas
    return t


def parse_data_type(text: str) -> tuple[dict, bool]:
    """OCR text -> ({"data_type","content"}, parsed_ok). Deterministic parse with light
    repair; content is a compact JSON string of the parsed structure. Returns
    parsed_ok=False when nothing parsed (caller decides on a VLM fallback)."""
    dt = detect_data_type(text)["data_type"]
    raw = _repair_json(text or "")
    parsed = None
    try:
        parsed = json.loads(raw)
    except Exception:
        try:
            import yaml
            y = yaml.safe_load(raw)
            if isinstance(y, (dict, list)):
                parsed = y
        except Exception:
            parsed = None
    if parsed is None:
        return {"data_type": dt, "content": ""}, False
    return {"data_type": dt, "content": json.dumps(parsed, ensure_ascii=False, separators=(",", ":"))}, True