qwen_test_runner/vision/fuse_schema.py

"""
fuse_schema.py — the FusedScene pydantic schema (fusion tier).

Hand-written nested models (the SubjectValue precedent — the SlotSpec codegen in
schema.py handles one nesting level; FusedScene needs three: entities ->
attributes -> ownership). Deliberately NOT a registered VisionTaskSpec: FusedScene
is a deterministically-produced dataset artifact, not a VLM probe — it needs no
GBNF, no system prompt, no scorer. If a VLM is ever trained to EMIT FusedScene,
register a flattened variant then.

Versioned from day one: this module is a second schema authority next to the
registry, so every instance stamps `fused_scene_version`.

Coordinate policy: all geometry emitted in the declared `coord_space`
(NORM_0_1000 by default, via specialists.box_to_space/poly_to_space); the one
pixel-unit field is `image_size`, documented as such. Scorers convert via
coords.to_canonical.
"""

from __future__ import annotations

from typing import Optional

from pydantic import BaseModel, Field

FUSED_SCENE_VERSION = "1.0"

# Caps (data, referenced by fuse.py — never hardcoded at call sites)
MAX_ENTITIES = 12            # kept by saliency rank
MAX_RELATION_ENTITIES = 8    # pairwise relations among the top-K entities
MASK_POLY_MAX_POINTS = 64    # outline_polygon(max_points=...) per entity

# Basin reasons (closed set, tested)
BASIN_REASONS = ("low_margin", "no_grounding_multi_entity", "ambiguous_binding",
                 "abstract_unbound")

# Ownership methods (closed set, tested)
OWNERSHIP_METHODS = ("single_entity", "mask_containment", "box_containment",
                     "caption_binding")


class GridPosition(BaseModel):
    grid: str                                   # "{upper|middle|lower} {left|center|right}"
    offset_from_center: list[float]             # (centroid-center)/half-extents, in [-1,1]


class DepthInfo(BaseModel):
    nearness: float                             # continuous [0,1], bigger = nearer
    rank: int                                   # 1 = nearest


class SaliencyInfo(BaseModel):
    score: float
    rank: int                                   # 1 = most salient


class MaskInfo(BaseModel):
    polygon: list[float] = Field(default_factory=list)   # flat x,y interleaved, task space
    quality: Optional[float] = None             # SAM predicted-IoU (retained signal)


class CaptionBinding(BaseModel):
    source: str                                 # caption column key
    subject_name: str
    confidence: float


class Ownership(BaseModel):
    confidence: float
    margin: Optional[float] = None              # f1 - f2 (containment methods only)
    method: str                                 # one of OWNERSHIP_METHODS


class RegionOnOwner(BaseModel):
    vertical: str                               # upper | middle | lower
    horizontal: str                             # left | center | right
    offset: list[float]                         # (attr center - owner centroid)/owner half-extents


class AttributeRecord(BaseModel):
    text: str                                   # canonical (longest) form after dedup
    stratum: str
    sources: list[str]                          # caption columns that carried it
    consensus: float                            # len(sources)/n_caption_sources
    grounded: bool = False
    box: Optional[list[float]] = None           # present iff grounded (task space)
    grounding_score: Optional[float] = None     # GDINO phrase score
    ownership: Ownership
    region_on_owner: Optional[RegionOnOwner] = None


class Entity(BaseModel):
    id: str                                     # person_1, person_2, dog, ... (left-to-right)
    label: str
    detection_score: float
    box: list[float]                            # task space
    centroid: list[float]                       # task space
    area_frac: float
    position: GridPosition
    depth: Optional[DepthInfo] = None
    saliency: SaliencyInfo
    is_primary: bool = False
    mask: Optional[MaskInfo] = None
    caption_bindings: list[CaptionBinding] = Field(default_factory=list)
    attributes: list[AttributeRecord] = Field(default_factory=list)


class Relation(BaseModel):
    a: str                                      # entity id (smaller entity index)
    b: str
    predicates: list[str]                       # spatial_relations predicate vocab
    dx: float                                   # (centroid_b - centroid_a)/W, task-space-free
    dy: float
    distance: float                             # centroid distance / image diagonal
    iou: float
    depth_delta: Optional[float] = None         # nearness_a - nearness_b (continuous)
    confidence: float                           # min(detection scores)


class BasinCandidate(BaseModel):
    entity_id: str
    likelihood: float


class BasinItem(BaseModel):
    text: str
    stratum: str
    sources: list[str]
    consensus: float
    reason: str                                 # one of BASIN_REASONS
    grounded: bool = False
    box: Optional[list[float]] = None
    candidates: list[BasinCandidate] = Field(default_factory=list)


class VotedValue(BaseModel):
    value: Optional[str] = None
    votes: dict[str, int] = Field(default_factory=dict)


class StyleValue(BaseModel):
    value: Optional[str] = None                 # specialist wins conflicts (it saw the image)
    caption_votes: dict[str, int] = Field(default_factory=dict)
    specialist: Optional[str] = None


class MoodValue(BaseModel):
    value: Optional[str] = None
    per_source: dict[str, str] = Field(default_factory=dict)


class SymmetryInfo(BaseModel):
    axis: str = "none"
    lr: Optional[float] = None                  # continuous correlations (retained)
    tb: Optional[float] = None


class SceneAttribute(BaseModel):
    text: str
    stratum: str = "scene_level"
    sources: list[str] = Field(default_factory=list)


class SceneOCRLine(BaseModel):
    text: str
    box: Optional[list[float]] = None
    conf: Optional[float] = None                # EasyOCR confidence (retained)


class SceneOCR(BaseModel):
    full_text: str = ""
    lines: list[SceneOCRLine] = Field(default_factory=list)


class LabelScore(BaseModel):
    label: str
    score: float


class SceneBlock(BaseModel):
    setting: VotedValue = Field(default_factory=VotedValue)
    style: StyleValue = Field(default_factory=StyleValue)
    mood: MoodValue = Field(default_factory=MoodValue)
    layout: str = "unknown"
    symmetry: SymmetryInfo = Field(default_factory=SymmetryInfo)
    actions: list[SceneAttribute] = Field(default_factory=list)   # unbound caption actions
    scene_attributes: list[SceneAttribute] = Field(default_factory=list)
    ocr: SceneOCR = Field(default_factory=SceneOCR)
    class_top: list[LabelScore] = Field(default_factory=list)


class Counts(BaseModel):
    total_entities: int = 0
    people: int = 0                             # via the person synonym group
    by_label: dict[str, int] = Field(default_factory=dict)


class GroundingStats(BaseModel):
    phrases_total: int = 0
    phrases_grounded: int = 0
    assigned: int = 0
    basin: int = 0
    scene_level: int = 0
    ungroundable: int = 0


class Quality(BaseModel):
    n_caption_sources: int = 0
    detection_score_mean: float = 0.0
    mask_quality_mean: Optional[float] = None
    ocr_conf_mean: Optional[float] = None
    grounding: GroundingStats = Field(default_factory=GroundingStats)
    overall_confidence: float = 0.0             # the fusion_confidence scalar


class FusedScene(BaseModel):
    fused_scene_version: str = FUSED_SCENE_VERSION
    coord_space: str
    image_size: list[int]                       # (W, H) PIXELS — the one pixel field
    counts: Counts = Field(default_factory=Counts)
    entities: list[Entity] = Field(default_factory=list)
    relations: list[Relation] = Field(default_factory=list)
    shared_basin: list[BasinItem] = Field(default_factory=list)
    scene: SceneBlock = Field(default_factory=SceneBlock)
    quality: Quality = Field(default_factory=Quality)


FUSED_SCENE_JSON_SCHEMA = FusedScene.model_json_schema()