coco_classes.py

from __future__ import annotations

import difflib
import functools
import logging
import re
from typing import Dict, Optional, Tuple

import numpy as np

logger = logging.getLogger(__name__)

COCO_CLASSES: Tuple[str, ...] = (
    "person",
    "bicycle",
    "car",
    "motorcycle",
    "airplane",
    "bus",
    "train",
    "truck",
    "boat",
    "traffic light",
    "fire hydrant",
    "stop sign",
    "parking meter",
    "bench",
    "bird",
    "cat",
    "dog",
    "horse",
    "sheep",
    "cow",
    "elephant",
    "bear",
    "zebra",
    "giraffe",
    "backpack",
    "umbrella",
    "handbag",
    "tie",
    "suitcase",
    "frisbee",
    "skis",
    "snowboard",
    "sports ball",
    "kite",
    "baseball bat",
    "baseball glove",
    "skateboard",
    "surfboard",
    "tennis racket",
    "bottle",
    "wine glass",
    "cup",
    "fork",
    "knife",
    "spoon",
    "bowl",
    "banana",
    "apple",
    "sandwich",
    "orange",
    "broccoli",
    "carrot",
    "hot dog",
    "pizza",
    "donut",
    "cake",
    "chair",
    "couch",
    "potted plant",
    "bed",
    "dining table",
    "toilet",
    "tv",
    "laptop",
    "mouse",
    "remote",
    "keyboard",
    "cell phone",
    "microwave",
    "oven",
    "toaster",
    "sink",
    "refrigerator",
    "book",
    "clock",
    "vase",
    "scissors",
    "teddy bear",
    "hair drier",
    "toothbrush",
)


def coco_class_catalog() -> str:
    """Return the COCO classes in a comma-separated catalog for prompts."""

    return ", ".join(COCO_CLASSES)


def _normalize(label: str) -> str:
    return re.sub(r"[^a-z0-9]+", " ", label.lower()).strip()


_CANONICAL_LOOKUP: Dict[str, str] = {_normalize(name): name for name in COCO_CLASSES}
_COCO_SYNONYMS: Dict[str, str] = {
    "people": "person",
    "man": "person",
    "woman": "person",
    "men": "person",
    "women": "person",
    "pedestrian": "person",
    "soldier": "person",
    "infantry": "person",
    "civilian": "person",
    "motorbike": "motorcycle",
    "motor bike": "motorcycle",
    "bike": "bicycle",
    "aircraft": "airplane",
    "plane": "airplane",
    "jet": "airplane",
    "aeroplane": "airplane",
    "drone": "airplane",
    "uav": "airplane",
    "helicopter": "airplane",
    "pickup": "truck",
    "pickup truck": "truck",
    "semi": "truck",
    "lorry": "truck",
    "tractor trailer": "truck",
    "vehicle": "car",
    "sedan": "car",
    "suv": "car",
    "van": "car",
    "vessel": "boat",
    "ship": "boat",
    "warship": "boat",
    "speedboat": "boat",
    "cargo ship": "boat",
    "fishing boat": "boat",
    "yacht": "boat",
    "kayak": "boat",
    "canoe": "boat",
    "watercraft": "boat",
    "coach": "bus",
    "television": "tv",
    "tv monitor": "tv",
    "mobile phone": "cell phone",
    "smartphone": "cell phone",
    "cellphone": "cell phone",
    "dinner table": "dining table",
    "sofa": "couch",
    "cooker": "oven",
}
_ALIAS_LOOKUP: Dict[str, str] = {_normalize(alias): canonical for alias, canonical in _COCO_SYNONYMS.items()}


# ---------------------------------------------------------------------------
# Semantic similarity fallback (lazy-loaded)
# ---------------------------------------------------------------------------

_SEMANTIC_MODEL = None
_COCO_EMBEDDINGS: Optional[np.ndarray] = None
_SEMANTIC_THRESHOLD = 0.65  # Minimum cosine similarity to accept a match


def _get_semantic_model():
    """Lazy-load a lightweight sentence-transformer for semantic matching."""
    global _SEMANTIC_MODEL, _COCO_EMBEDDINGS
    if _SEMANTIC_MODEL is not None:
        return _SEMANTIC_MODEL, _COCO_EMBEDDINGS

    try:
        from sentence_transformers import SentenceTransformer
        _SEMANTIC_MODEL = SentenceTransformer("all-MiniLM-L6-v2")
        # Prefix with "a photo of a" to anchor embeddings in visual/object space
        coco_phrases = [f"a photo of a {cls}" for cls in COCO_CLASSES]
        _COCO_EMBEDDINGS = _SEMANTIC_MODEL.encode(
            coco_phrases, normalize_embeddings=True
        )
        logger.info("Loaded semantic similarity model for COCO class mapping")
    except Exception:
        logger.warning("sentence-transformers unavailable; semantic COCO mapping disabled", exc_info=True)
        _SEMANTIC_MODEL = False  # Sentinel: tried and failed
        _COCO_EMBEDDINGS = None

    return _SEMANTIC_MODEL, _COCO_EMBEDDINGS


def _semantic_coco_match(value: str) -> Optional[str]:
    """Find the closest COCO class by embedding cosine similarity.

    Returns the COCO class name if similarity >= threshold, else None.
    """
    model, coco_embs = _get_semantic_model()
    if model is False or coco_embs is None:
        return None

    query_emb = model.encode(
        [f"a photo of a {value}"], normalize_embeddings=True
    )
    similarities = query_emb @ coco_embs.T  # (1, 80)
    best_idx = int(np.argmax(similarities))
    best_score = float(similarities[0, best_idx])

    if best_score >= _SEMANTIC_THRESHOLD:
        matched = COCO_CLASSES[best_idx]
        logger.info(
            "Semantic COCO match: '%s' -> '%s' (score=%.3f)",
            value, matched, best_score,
        )
        return matched

    logger.debug(
        "Semantic COCO match failed: '%s' best='%s' (score=%.3f < %.2f)",
        value, COCO_CLASSES[best_idx], best_score, _SEMANTIC_THRESHOLD,
    )
    return None


@functools.lru_cache(maxsize=512)
def canonicalize_coco_name(value: str | None) -> str | None:
    """Map an arbitrary string to the closest COCO class name if possible.

    Matching cascade:
    1. Exact normalized match
    2. Synonym lookup
    3. Substring match (alias then canonical)
    4. Token-level match
    5. Fuzzy string match (difflib)
    6. Semantic embedding similarity (sentence-transformers)
    """

    if not value:
        return None
    normalized = _normalize(value)
    if not normalized:
        return None
    if normalized in _CANONICAL_LOOKUP:
        return _CANONICAL_LOOKUP[normalized]
    if normalized in _ALIAS_LOOKUP:
        return _ALIAS_LOOKUP[normalized]

    for alias_norm, canonical in _ALIAS_LOOKUP.items():
        if alias_norm and alias_norm in normalized:
            return canonical
    for canonical_norm, canonical in _CANONICAL_LOOKUP.items():
        if canonical_norm and canonical_norm in normalized:
            return canonical

    tokens = normalized.split()
    for token in tokens:
        if token in _CANONICAL_LOOKUP:
            return _CANONICAL_LOOKUP[token]
        if token in _ALIAS_LOOKUP:
            return _ALIAS_LOOKUP[token]

    close = difflib.get_close_matches(normalized, list(_CANONICAL_LOOKUP.keys()), n=1, cutoff=0.82)
    if close:
        return _CANONICAL_LOOKUP[close[0]]

    # Last resort: semantic embedding similarity
    return _semantic_coco_match(value)