qwen_test_runner/vision/strata.py

"""
strata.py — the attribute stratification lexicon (fusion tier).

Classifies a caption attribute string into a STRATUM — the "disperse the topics"
layer that replaces the flat attribute list with typed, routable records. Pure
stdlib, deterministic, registry-as-python (same pattern as registry.py /
tasks_vision.py): the lexicon is data to iterate on, not code.

Routing semantics consumed by fuse.py:
  - GROUNDABLE strata are sent to the GDINO phrase-grounding pass (they name
    visible things a detector can box).
  - "scene_level" bypasses entities entirely -> FusedScene.scene.scene_attributes.
  - "abstract_quality", "color", and "action" ride the caption-binding-only path
    (never grounded: a detector box for "elegant" or bare "red" is noise).
  - Everything is classified; "abstract_quality" is the catch-all default.
"""

from __future__ import annotations

import re

from .metrics import _depluralize

# Minimal stopword set for head-noun extraction (articles/preps/conjunctions that
# can trail a phrase). Deliberately tiny — attribute phrases are short.
_STOP = frozenset({
    "a", "an", "the", "of", "in", "on", "at", "with", "and", "or", "to",
    "her", "his", "its", "their", "very", "slightly",
})

_TOKEN_RE = re.compile(r"[a-z0-9]+(?:-[a-z0-9]+)*")


STRATA: dict[str, frozenset] = {
    "hair": frozenset({
        "hair", "hairstyle", "bangs", "fringe", "ponytail", "pigtails", "twintails",
        "braid", "braids", "bun", "curls", "updo", "bob", "undercut", "mohawk",
        "sidelocks", "ahoge", "afro", "dreadlocks", "cornrows", "mullet", "buzzcut",
    }),
    "face": frozenset({
        "face", "eyes", "eye", "eyebrows", "eyebrow", "eyelashes", "lips", "lip",
        "mouth", "nose", "cheeks", "cheekbones", "chin", "jaw", "jawline", "forehead",
        "freckles", "dimples", "beard", "mustache", "stubble", "smile", "grin",
        "expression", "gaze", "makeup", "lipstick", "eyeliner", "eyeshadow", "blush",
        "mascara", "teeth",
    }),
    "skin": frozenset({
        "skin", "complexion", "tan", "tattoo", "tattoos", "scar", "scars", "mole",
        "birthmark", "wrinkles", "pores",
    }),
    "clothing": frozenset({
        "dress", "shirt", "t-shirt", "tshirt", "blouse", "top", "skirt", "pants",
        "trousers", "jeans", "shorts", "jacket", "coat", "hoodie", "sweater",
        "cardigan", "vest", "suit", "uniform", "kimono", "yukata", "robe", "gown",
        "leotard", "swimsuit", "bikini", "armor", "cape", "cloak", "apron",
        "sleeves", "sleeve", "collar", "neckline", "hem", "outfit", "attire",
        "clothes", "clothing", "costume", "sweatshirt", "leggings", "stockings",
        "tights", "socks", "corset", "bodysuit", "tunic", "sari", "poncho",
    }),
    "accessory": frozenset({
        "earrings", "earring", "necklace", "pendant", "choker", "bracelet", "ring",
        "rings", "watch", "hat", "cap", "beanie", "beret", "crown", "tiara",
        "headband", "hairband", "ribbon", "bow", "hairpin", "hairclip", "scrunchie",
        "glasses", "sunglasses", "eyepatch", "monocle", "mask", "scarf", "gloves",
        "glove", "belt", "bag", "handbag", "backpack", "purse", "umbrella", "fan",
        "brooch", "badge", "piercing", "anklet", "shoes", "boots", "sandals",
        "heels", "sneakers", "veil", "headphones", "tie", "bowtie",
    }),
    "body": frozenset({
        "build", "figure", "physique", "body", "shoulders", "shoulder", "arms",
        "arm", "hands", "hand", "fingers", "legs", "leg", "thighs", "knees",
        "feet", "chest", "waist", "hips", "back", "neck", "collarbone", "height",
        "frame", "posture", "muscles", "abs", "curves",
    }),
    "pose": frozenset({
        "standing", "sitting", "kneeling", "crouching", "lying", "leaning",
        "walking", "running", "jumping", "dancing", "posing", "looking", "facing",
        "reaching", "pointing", "waving", "holding", "carrying", "crossed",
        "outstretched", "tilted", "turned", "pose", "stance",
    }),
    "color": frozenset({
        "red", "orange", "yellow", "green", "blue", "purple", "violet", "pink",
        "brown", "black", "white", "gray", "grey", "silver", "gold", "golden",
        "blonde", "blond", "brunette", "auburn", "crimson", "scarlet", "teal",
        "turquoise", "cyan", "magenta", "lavender", "beige", "cream", "ivory",
        "navy", "maroon", "olive", "platinum", "pastel", "neon", "dark", "light",
        "pale", "bright", "vivid", "striped", "plaid", "polka-dot", "checkered",
        "floral", "gradient",
    }),
    "abstract_quality": frozenset({
        "beautiful", "pretty", "handsome", "cute", "elegant", "graceful", "stylish",
        "fashionable", "detailed", "intricate", "delicate", "soft", "sharp",
        "masterpiece", "quality", "aesthetic", "gorgeous", "stunning", "charming",
        "youthful", "mature", "young", "old", "confident", "shy", "serene", "calm",
        "cheerful", "melancholic", "mysterious", "dramatic", "ethereal", "dreamy",
    }),
    "scene_level": frozenset({
        "background", "foreground", "backdrop", "lighting", "light", "shadow",
        "shadows", "sunlight", "moonlight", "sunset", "sunrise", "dusk", "dawn",
        "sky", "clouds", "bokeh", "blur", "depth", "wall", "walls", "floor",
        "ceiling", "window", "windows", "door", "room", "indoors", "outdoors",
        "outdoor", "indoor", "scenery", "landscape", "cityscape", "street",
        "forest", "beach", "mountains", "atmosphere", "ambiance", "setting",
        "scene", "environment", "composition", "framing",
    }),
}

# hyphen/compound-adjective suffixes -> stratum ("silver-haired", "blue-eyed")
SUFFIX_RULES: tuple = (
    ("haired", "hair"),
    ("eyed", "face"),
    ("faced", "face"),
    ("skinned", "skin"),
    ("sleeved", "clothing"),
    ("dressed", "clothing"),
    ("clad", "clothing"),
    ("shouldered", "body"),
    ("legged", "body"),
    ("armed", "body"),
)

# -ing words that are NOUNS, not gerunds — exempt from the verb-phrase rule
# (data to extend as COCO round-trips surface more)
_NOUN_ING = frozenset({
    "wedding", "building", "painting", "lighting", "ceiling", "clothing",
    "evening", "morning", "string", "earring", "ring", "king", "wing",
    "railing", "awning",
})

# Strata whose phrases go to the GDINO grounding pass (visible, boxable things).
GROUNDABLE = frozenset({"hair", "face", "skin", "clothing", "accessory", "body", "pose"})

# Any-token tie-break order (only reached when the head noun missed the lexicon).
# Concrete/visible strata outrank colors and abstractions.
STRATUM_PRECEDENCE = ("hair", "face", "skin", "accessory", "clothing", "body",
                      "pose", "scene_level", "color", "abstract_quality")

# The full stratum vocabulary fuse.py may emit ("action" is assigned by fuse.py to
# caption `actions` entries directly — it has no lexicon and is never grounded).
ALL_STRATA = tuple(STRATA.keys()) + ("action",)


def _content_tokens(text: str) -> list:
    return [t for t in _TOKEN_RE.findall((text or "").lower()) if t not in _STOP]


def classify_stratum(text: str) -> str:
    """Deterministic stratum for an attribute string.

    1. head-noun rule: depluralized LAST content token, exact lexicon lookup
    2. suffix rules on the head token ("silver-haired" -> hair)
    3. any-token lookup in STRATUM_PRECEDENCE order
    4. all tokens are color/pattern terms -> color
    5. default -> abstract_quality (nothing is ever unclassified)
    """
    toks = _content_tokens(text)
    if not toks:
        return "abstract_quality"
    # _depluralize is crude ("dress"->"dres") — always try the raw form too
    head_forms = {toks[-1], _depluralize(toks[-1])}

    for stratum, words in STRATA.items():
        if head_forms & words:
            return stratum
    for suffix, stratum in SUFFIX_RULES:
        if any(h.endswith(suffix) for h in head_forms):
            return stratum
    forms = [{t, _depluralize(t)} for t in toks]
    for stratum in STRATUM_PRECEDENCE:
        words = STRATA[stratum]
        if any(f & words for f in forms):
            return stratum
    if all(f & STRATA["color"] for f in forms):
        return "color"
    # verb-phrase heuristic: leading gerund ("playing baseball", "taking a photo",
    # "running") → pose. Caught live on COCO captions, where the structurer emits
    # verb phrases as attributes that otherwise fell to abstract_quality.
    first = toks[0]
    if first.endswith("ing") and len(first) > 4 and first not in _NOUN_ING:
        return "pose"
    return "abstract_quality"


def is_groundable(stratum: str) -> bool:
    return stratum in GROUNDABLE