""" 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": "", "confidence": , ' '"top5": [{"label": "", "score": }]}' ), 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": "", "box": [x1, y1, x2, y2], "score": }], ' '"count": }\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": "", ' '"lines": [{"text": "", "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": "", ' '"confidence": }' ), 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": "", "predicate": ' '"", "object": ""}]}' ), 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\": \"\", \"farthest\": \"\", " '"relative_depth": [{"a": "", "b": "", "a_is": ""}]}' ), 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": "", "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": "", "content": ""}' ), 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": "", "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": ""}\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": "", "bbox3d": [x, y, z, w, h, l, yaw], ' '"score": }]}\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\": [, , ]}\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": "", "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": "", "relation": ' '"", "b": ""}]}' ), 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": "", ' '"layout": "", ' '"symmetry": ""}' ), 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