qwen_test_runner/vision/tasks_vision.py

"""
tasks_vision.py — The 15 vision-task categories, as data.

Each VisionTaskSpec owns a small per-category field registry (dict[str, SlotSpec])
and a system/user prompt. The Pydantic model, JSON Schema, GBNF grammar, and
Claude tool schema are generated from that registry by the SAME machinery the
caption schema uses (schema.build_*). Adding a category is one dict entry.

Three categories are PILOT (full schema + GT dataset + real metric); the other
twelve are STUB (valid minimal schema so their grammar builds, metric wired in
Phase 3). This mirrors how registry.py grows the caption schema.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Mapping

from ..registry import SlotSpec
from ..schema import build_gbnf_from_registry, build_json_schema, build_model_from_registry
from .coords import CoordSpace


# ──────────────────────────────────────────────────────────────────────────────
# Field-builder shorthand (keeps the registry readable)
# ──────────────────────────────────────────────────────────────────────────────

def _f(name, **kw) -> SlotSpec:
    """A single-value open string field unless overridden."""
    kw.setdefault("cardinality", "single")
    kw.setdefault("vocabulary", "open")
    return SlotSpec(name=name, **kw)


def _enum(name, values, optional=False) -> SlotSpec:
    return SlotSpec(name=name, cardinality="single", vocabulary="closed",
                    closed_values=tuple(values), optional=optional)


def _list_of(name, *fields, max_items=32) -> SlotSpec:
    return SlotSpec(name=name, cardinality="list", vocabulary="open",
                    nested_fields=tuple(fields), max_items=max_items)


@dataclass(frozen=True)
class VisionTaskSpec:
    category: str
    probes: str
    fields: Mapping[str, SlotSpec]
    system_prompt: str
    user_prompt: str
    metric: str                       # key into metrics._SCORERS
    status: str = "pilot"             # "pilot" | "stub"
    coord_space: CoordSpace = CoordSpace.NORM_0_1000
    gt_dataset: str = ""              # key into datasets.DATASET_REGISTRY
    gt_split: str = ""
    max_new_tokens: int = 512
    license_note: str = ""
    download_gb: float = 0.0
    per_sample_prompt: bool = False   # use GTSample.prompt as the user prompt (e.g. VQA question)


# Generated-artifact caches (keyed by category — VisionTaskSpec holds a dict so
# it isn't hashable; categories are unique).
_MODEL_CACHE: dict[str, type] = {}
_GBNF_CACHE: dict[str, str] = {}


def model_for(spec: VisionTaskSpec):
    if spec.category not in _MODEL_CACHE:
        _MODEL_CACHE[spec.category] = build_model_from_registry(
            "Vision_" + spec.category.title().replace("_", ""), dict(spec.fields)
        )
    return _MODEL_CACHE[spec.category]


def json_schema_for(spec: VisionTaskSpec) -> dict:
    return build_json_schema(model_for(spec))


def gbnf_for(spec: VisionTaskSpec) -> str:
    if spec.category not in _GBNF_CACHE:
        _GBNF_CACHE[spec.category] = build_gbnf_from_registry(dict(spec.fields))
    return _GBNF_CACHE[spec.category]


def tool_schema_for(spec: VisionTaskSpec) -> dict:
    """Claude-style tool input_schema (the per-category JSON Schema)."""
    return json_schema_for(spec)


# ──────────────────────────────────────────────────────────────────────────────
# PILOT categories (full)
# ──────────────────────────────────────────────────────────────────────────────

_CLASSIFICATION = VisionTaskSpec(
    category="image_classification",
    probes="native ViT classification emitted as JSON",
    fields={
        "label": _f("label", optional=False, max_str_length=64),
        "confidence": _f("confidence", value_kind="number", optional=False, number_range=(0.0, 1.0)),
        "top5": _list_of(
            "top5",
            _f("label", optional=False, max_str_length=64),
            _f("score", value_kind="number", optional=False, number_range=(0.0, 1.0)),
            max_items=5,
        ),
    },
    system_prompt=(
        "You are an image classifier. Identify the single most prominent object or scene "
        "category in the image. Output ONLY a raw JSON object and NOTHING else — no prose, "
        "no explanation, and NO markdown code fences (do not wrap it in ```). "
        "It must match this shape exactly:\n"
        '{"label": "<string>", "confidence": <number 0..1>, '
        '"top5": [{"label": "<string>", "score": <number 0..1>}]}'
    ),
    user_prompt="Classify this image. Output only the raw JSON object.",
    metric="classification",
    gt_dataset="imagenet_val",
    gt_split="validation",
    max_new_tokens=160,
    license_note="ImageNet: non-commercial research use.",
)

_BBOX = VisionTaskSpec(
    category="bbox_grounding",
    probes="object localization + grounded counting",
    fields={
        "detections": _list_of(
            "detections",
            _f("label", optional=False, max_str_length=64),
            _f("box", value_kind="bbox", optional=False),
            _f("score", value_kind="number", optional=False, number_range=(0.0, 1.0)),
            max_items=32,
        ),
        "count": _f("count", value_kind="integer", optional=False),
    },
    system_prompt=(
        "You are an object detector. Find every distinct object in the image. Output ONLY a "
        "raw JSON object and NOTHING else — no prose, no markdown code fences (do not wrap it "
        "in ```). It must match this shape exactly:\n"
        '{"detections": [{"label": "<string>", "box": [x1, y1, x2, y2], "score": <number 0..1>}], '
        '"count": <integer>}\n'
        "{coord_hint} Use the key \"box\" (NOT bbox_2d) with exactly four numbers [x1, y1, x2, y2]."
    ),
    user_prompt="Detect all objects in this image. Output only the raw JSON object.",
    metric="detection",
    coord_space=CoordSpace.NORM_0_1000,
    gt_dataset="coco_detection",
    gt_split="val",
    max_new_tokens=768,
    license_note="COCO: CC-BY 4.0 (images vary).",
)

_OCR = VisionTaskSpec(
    category="ocr_text",
    probes="text reading + transcription fidelity + localization",
    fields={
        "full_text": _f("full_text", optional=False, max_str_length=4096),
        "lines": _list_of(
            "lines",
            _f("text", optional=False, max_str_length=512),
            _f("box", value_kind="bbox", optional=True),
            max_items=64,
        ),
    },
    system_prompt=(
        "You are an OCR engine. Transcribe all readable text in the image. Output ONLY a raw "
        "JSON object and NOTHING else — no prose, no markdown code fences (do not wrap it in "
        "```). It must match this shape exactly:\n"
        '{"full_text": "<all text, joined by spaces>", '
        '"lines": [{"text": "<string>", "box": [x1, y1, x2, y2]}]}\n'
        "{coord_hint} If you cannot localize a line, omit its box."
    ),
    user_prompt="Read all the text in this image. Output only the raw JSON object.",
    metric="ocr",
    coord_space=CoordSpace.NORM_0_1000,
    gt_dataset="textvqa",
    gt_split="validation",
    max_new_tokens=512,
    license_note="TextVQA: CC-BY 4.0.",
)


# ──────────────────────────────────────────────────────────────────────────────
# STUB categories (minimal valid schema; metric + GT wired in Phase 3)
# ──────────────────────────────────────────────────────────────────────────────

def _stub(category, probes, fields, prompt, **kw) -> VisionTaskSpec:
    kw.setdefault("metric", "schema_only")
    kw.setdefault("status", "stub")
    kw.setdefault("user_prompt", "Analyze this image.")
    return VisionTaskSpec(category=category, probes=probes, fields=fields,
                          system_prompt=prompt, **kw)


_SPATIAL_PREDS = ("left_of", "right_of", "above", "below", "on", "under",
                  "inside", "behind", "in_front_of")

_STUBS = []

_DATATYPE_VALUES = ("json", "yaml", "markdown", "csv", "toml", "xml", "code", "plaintext")

_DATATYPE_DIFF = VisionTaskSpec(
    category="data_type_differentiation",
    probes="recognize a rendered data format from a screenshot",
    fields={
        "data_type": _enum("data_type", _DATATYPE_VALUES),
        "confidence": _f("confidence", value_kind="number", optional=False, number_range=(0.0, 1.0)),
    },
    system_prompt=(
        "You are shown a screenshot of structured data. Identify which serialization format "
        "it is. Output ONLY a raw JSON object, no markdown fences:\n"
        '{"data_type": "<one of: json, yaml, markdown, csv, toml, xml, code, plaintext>", '
        '"confidence": <number 0..1>}'
    ),
    user_prompt="What data format is shown? Output only the raw JSON object.",
    metric="datatype_diff",
    gt_dataset="datatype_synth",
    max_new_tokens=96,
    license_note="synthetic (self-contained).",
)

_SPATIAL = VisionTaskSpec(
    category="structural_spatial_awareness",
    probes="spatial relations between objects",
    fields={"relations": _list_of(
        "relations",
        _f("subject", optional=False),
        _enum("predicate", _SPATIAL_PREDS),
        _f("object", optional=False), max_items=12)},
    system_prompt=(
        "Describe the spatial relations between the colored shapes. Subjects and objects are "
        "the colors (red, green, blue). Output ONLY raw JSON, no fences:\n"
        '{"relations": [{"subject": "<color>", "predicate": '
        '"<left_of|right_of|above|below>", "object": "<color>"}]}'
    ),
    user_prompt="List the spatial relations between the colored shapes. Raw JSON only.",
    metric="triples", gt_dataset="shapes_synth", max_new_tokens=256,
    license_note="synthetic (self-contained).",
)

_DEPTH = VisionTaskSpec(
    category="depth_analysis",
    probes="relative depth ordering",
    fields={
        "nearest": _f("nearest"),
        "farthest": _f("farthest"),
        "relative_depth": _list_of(
            "relative_depth",
            _f("a", optional=False),
            _f("b", optional=False),
            _enum("a_is", ("nearer", "farther", "same")), max_items=12),
    },
    system_prompt=(
        "Judge relative depth of the colored shapes: a LARGER shape appears NEARER. Output ONLY "
        "raw JSON, no fences:\n{\"nearest\": \"<color>\", \"farthest\": \"<color>\", "
        '"relative_depth": [{"a": "<color>", "b": "<color>", "a_is": "<nearer|farther|same>"}]}'
    ),
    user_prompt="Report the relative depth of the colored shapes. Raw JSON only.",
    metric="depth_order", gt_dataset="shapes_synth", max_new_tokens=256,
    license_note="synthetic (self-contained).",
)

_SUBJECT = VisionTaskSpec(
    category="subject_fixation",
    probes="primary salient subject",
    fields={"primary_subject": SlotSpec(
        name="primary_subject", cardinality="single", vocabulary="open", optional=False,
        nested_fields=(_f("label", optional=False), _f("box", value_kind="bbox", optional=False)))},
    system_prompt=(
        "Identify the single most prominent (largest) shape — its color and bounding box. "
        "Output ONLY raw JSON, no fences:\n"
        '{"primary_subject": {"label": "<color>", "box": [x1, y1, x2, y2]}}\n{coord_hint}'
    ),
    user_prompt="Identify the primary subject and its box. Raw JSON only.",
    metric="subject_fixation", gt_dataset="shapes_synth", coord_space=CoordSpace.NORM_0_1000,
    max_new_tokens=128, license_note="synthetic (self-contained).",
)

_DATATYPE_UTIL = VisionTaskSpec(
    category="data_type_utilization",
    probes="parse a rendered data format into normalized JSON",
    fields={
        "data_type": _enum("data_type", _DATATYPE_VALUES),
        "content": _f("content", optional=False, max_str_length=2048),
    },
    system_prompt=(
        "You are shown a screenshot of structured data. Read it and re-serialize its contents "
        "as JSON. Output ONLY a raw JSON object, no markdown fences:\n"
        '{"data_type": "<the format>", "content": "<the data as a JSON string, e.g. '
        '{\\"name\\": \\"Alice\\"}>"}'
    ),
    user_prompt="Read the data and output {data_type, content} as raw JSON.",
    metric="datatype_util",
    gt_dataset="datatype_synth",
    max_new_tokens=512,
    license_note="synthetic (self-contained).",
)


# ──────────────────────────────────────────────────────────────────────────────
# THE REGISTRY
# ──────────────────────────────────────────────────────────────────────────────

_SEGMENTATION = VisionTaskSpec(
    category="segmentation",
    probes="instance segmentation as labeled polygons",
    fields={
        "masks": _list_of(
            "masks",
            _f("label", optional=False, max_str_length=64),
            SlotSpec(name="polygon", cardinality="list", vocabulary="open",
                     value_kind="number", max_items=512, optional=False),
            max_items=32,
        ),
    },
    system_prompt=(
        "You are an instance segmenter. Trace the outline of every distinct object "
        "as a closed polygon. Output ONLY a raw JSON object and NOTHING else — no prose, "
        "no markdown code fences (do not wrap it in ```). It must match this shape exactly:\n"
        '{"masks": [{"label": "<string>", "polygon": [x1, y1, x2, y2, x3, y3, ...]}]}\n'
        "All x, y values are integers in 0..1000 relative to the image width and height. "
        "Each polygon is a FLAT list of alternating x, y vertices — a closed shape with at "
        "least 3 points / 6 numbers tracing the object boundary in order. This is a POLYGON, "
        "NOT a 4-number bounding box."
    ),
    user_prompt="Segment every object in this image as a labeled polygon. Output only the raw JSON object.",
    metric="segmentation",
    coord_space=CoordSpace.NORM_0_1000,
    gt_dataset="segmentation_synth",
    max_new_tokens=768,
    license_note="synthetic (self-contained).",
)

_OUTLINE = VisionTaskSpec(
    category="outline_association",
    probes="trace the main (largest) object's outline polygon + label it",
    fields={
        "outline": SlotSpec(name="outline", cardinality="list", vocabulary="open",
                            value_kind="number", max_items=256, optional=False),
        "label": _f("label", optional=False, max_str_length=64),
    },
    system_prompt=(
        "You are an outline tracer. Find the SINGLE largest (most prominent) object in the "
        "image and trace its outline as a closed polygon. Output ONLY a raw JSON object and "
        "NOTHING else - no prose, no markdown code fences (do not wrap it in ```). It must "
        "match this shape exactly:\n"
        '{"outline": [x1, y1, x2, y2, x3, y3, ...], "label": "<string>"}\n'
        "The outline is a flat list of alternating x, y vertex coordinates (at least 3 "
        "vertices = 6 numbers), tracing the object boundary in order. All x, y values are "
        "integers in 0..1000 relative to the image width and height. This is a POLYGON with "
        "MANY points, NOT a 4-number bounding box."
    ),
    user_prompt="Trace the main object's outline and label it. Output only the raw JSON object.",
    metric="outline_iou",
    status="pilot",
    coord_space=CoordSpace.NORM_0_1000,
    gt_dataset="outline_synth",
    max_new_tokens=640,
    license_note="synthetic (self-contained).",
)

_GEO3D = VisionTaskSpec(
    category="geometric_3d_object_id",
    probes="3D object identification with 3D boxes (simplified ground-plane proxy)",
    fields={
        "objects": _list_of(
            "objects",
            _f("class", optional=False, max_str_length=64),
            SlotSpec(name="bbox3d", cardinality="list", vocabulary="open",
                     value_kind="number", max_items=7, optional=False),
            _f("score", value_kind="number", optional=True, number_range=(0.0, 1.0)),
            max_items=16,
        ),
    },
    system_prompt=(
        "You are a 3D object detector looking at a scene of colored boxes resting on a "
        "ground plane. For each box report its class (its color) and a 3D bounding box. "
        "Output ONLY a raw JSON object and NOTHING else - no prose, no markdown code "
        "fences (do not wrap it in ```). It must match this shape exactly:\n"
        '{"objects": [{"class": "<color>", "bbox3d": [x, y, z, w, h, l, yaw], '
        '"score": <number 0..1>}]}\n'
        "All coordinates are normalized to 0..1 of the scene: x is the left-right ground "
        "position, z is the depth (0=near, 1=far), y is the height off the ground (0 on the "
        "floor); w, h, l are the box width, height and length; yaw is the rotation in "
        'radians. Use the key "bbox3d" with exactly seven numbers [x, y, z, w, h, l, yaw].'
    ),
    user_prompt="Identify the 3D boxes in this scene. Output only the raw JSON object.",
    metric="iou3d",
    status="pilot",
    coord_space=CoordSpace.NORM_0_1,
    gt_dataset="boxes3d_synth",
    max_new_tokens=384,
    license_note="synthetic (self-contained); simplified ground-plane 3D proxy.",
)

_CAMERA_ROT = VisionTaskSpec(
    category="camera_rotational_offset",
    probes="camera pose / rotation estimation from a 2D orientation cue",
    fields={
        "rotation": SlotSpec(name="rotation", cardinality="list", vocabulary="open",
                             value_kind="number", max_items=3, optional=False),
    },
    system_prompt=(
        "You estimate the camera's rotation relative to the scene. Output the three "
        "Euler angles in DEGREES as [yaw, pitch, roll]. Output ONLY a raw JSON object and "
        "NOTHING else — no prose, no explanation, and NO markdown code fences (do not wrap "
        "it in ```). It must match this shape exactly:\n"
        '{\"rotation\": [<yaw>, <pitch>, <roll>]}\n'
        "Each angle is a number in degrees in the range -180..180. If an axis is not "
        "discernible, report 0."
    ),
    user_prompt="Estimate the camera rotation [yaw, pitch, roll] in degrees. Output only the raw JSON object.",
    metric="angular_error",
    status="pilot",
    gt_dataset="camera_rot_synth",
    max_new_tokens=64,
    license_note="synthetic (self-contained).",
)

_VQA = VisionTaskSpec(
    category="vit_accuracy_to_prompt",
    probes="grounded visual question answering",
    fields={
        "answer": _f("answer", optional=False, max_str_length=512),
        "grounded_region": _f("grounded_region", value_kind="bbox", optional=True),
    },
    system_prompt=(
        "You are a visual question answering engine. Answer the user's question about "
        "the image as briefly as possible (a single word or short phrase). Optionally "
        "ground your answer with the bounding box of the region you used. Output ONLY a "
        "raw JSON object and NOTHING else — no prose, no explanation, and NO markdown "
        "code fences (do not wrap it in ```). It must match this shape exactly:\n"
        '{"answer": "<short answer>", "grounded_region": [x1, y1, x2, y2]}\n'
        "{coord_hint} If you cannot or need not localize, omit grounded_region entirely."
    ),
    user_prompt="Answer the question about this image. Output only the raw JSON object.",
    metric="vqa",
    per_sample_prompt=True,
    coord_space=CoordSpace.NORM_0_1000,
    gt_dataset="gqa",
    gt_split="validation",
    max_new_tokens=128,
    license_note="GQA / VQAv2: research use; images CC-BY (vary).",
)

_SEMANTIC = VisionTaskSpec(
    category="semantic_association",
    probes="semantic associations between entities as (a, relation, b) triples",
    fields={
        "associations": _list_of(
            "associations",
            _f("a", optional=False, max_str_length=64),
            _enum("relation", ("left_of", "right_of", "near", "is_a", "related_to")),
            _f("b", optional=False, max_str_length=64),
            max_items=32,
        ),
    },
    system_prompt=(
        "You relate the entities in the image to each other as semantic association "
        "triples. Each association links entity \"a\" to entity \"b\" by a relation. "
        "For the colored shapes, the entities are the colors (red, green, blue) and "
        "the shape type (circle). Allowed relations: left_of, right_of, near, is_a, "
        "related_to. Output ONLY a raw JSON object and NOTHING else - no prose, no "
        "explanation, and NO markdown code fences (do not wrap it in ```). It must "
        "match this shape exactly:\n"
        '{"associations": [{"a": "<entity>", "relation": '
        '"<left_of|right_of|near|is_a|related_to>", "b": "<entity>"}]}'
    ),
    user_prompt="List the semantic associations between the entities. Output only the raw JSON object.",
    metric="triples",
    gt_dataset="semantic_synth",
    max_new_tokens=384,
    license_note="synthetic (self-contained).",
)

_STYLE = VisionTaskSpec(
    category="style_structural_awareness",
    probes="visual style + structural layout/symmetry, as a coarse closed-vocab triple",
    fields={
        "style": _enum("style", ("photo", "painting", "3d_render", "sketch", "anime", "other")),
        "layout": _enum("layout", ("centered", "rule_of_thirds", "symmetric", "scattered", "unknown")),
        "symmetry": _enum("symmetry", ("horizontal", "vertical", "radial", "none")),
    },
    system_prompt=(
        "You judge the VISUAL STYLE and STRUCTURE of an image. Pick exactly one value "
        "from each closed vocabulary. Output ONLY a raw JSON object and NOTHING else — no "
        "prose, no explanation, and NO markdown code fences (do not wrap it in ```). "
        "It must match this shape exactly:\n"
        '{"style": "<one of: photo, painting, 3d_render, sketch, anime, other>", '
        '"layout": "<one of: centered, rule_of_thirds, symmetric, scattered, unknown>", '
        '"symmetry": "<one of: horizontal, vertical, radial, none>"}'
    ),
    user_prompt="Classify the visual style and structure. Output only the raw JSON object.",
    metric="style",
    status="pilot",
    gt_dataset="style_synth",
    max_new_tokens=96,
    license_note="synthetic (self-contained).",
)


VISION_TASK_REGISTRY: dict[str, VisionTaskSpec] = {
    t.category: t for t in [_CLASSIFICATION, _BBOX, _OCR, _DATATYPE_DIFF, _DATATYPE_UTIL,
                            _SPATIAL, _DEPTH, _SUBJECT,
                            _SEGMENTATION, _OUTLINE, _GEO3D, _CAMERA_ROT, _VQA, _SEMANTIC, _STYLE]
}


def get_task(category: str) -> VisionTaskSpec:
    if category not in VISION_TASK_REGISTRY:
        raise KeyError(f"unknown vision category: {category!r}. known: {list(VISION_TASK_REGISTRY)}")
    return VISION_TASK_REGISTRY[category]


def category_names() -> list[str]:
    return list(VISION_TASK_REGISTRY.keys())


def pilot_categories() -> list[str]:
    return [c for c, t in VISION_TASK_REGISTRY.items() if t.status == "pilot"]


def resolved_system_prompt(spec: VisionTaskSpec) -> str:
    """Fill the {coord_hint} placeholder using the task's coord_space."""
    if "{coord_hint}" in spec.system_prompt:
        from .coords import prompt_hint_for
        return spec.system_prompt.replace("{coord_hint}", prompt_hint_for(spec.coord_space))
    return spec.system_prompt