FabledCurator/backend/app/services/ml/ccip.py

"""CCIP few-shot character matcher (#114) — server-side, numpy on stored vectors.

CCIP is a FROZEN identity embedding; we don't train it. Instead the operator's
tagged characters become reference prototypes: a character tag's references are
the CCIP vectors of figure/face regions on images carrying that tag. To suggest
characters for a new image, we compare its figure-region CCIP vectors to every
character's references (multi-prototype: best match over a character's examples)
and surface the ones that clear a similarity threshold. No GPU here — the agent
already produced the vectors; this is cosine matching on what's stored.

v1 uses cosine similarity on the raw CCIP vectors with a tunable threshold; the
exact CCIP difference metric/threshold gets validated against the model during
the hands-on eval. numpy is imported lazily (API worker has it via pgvector).
"""

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import ImageRegion, MLSettings, Tag, TagKind
from ...models.tag import image_tag

# Cosine-similarity floor to call a figure the same character. The live setting
# (ml_settings.ccip_match_threshold) drives it; this is only the fallback when no
# threshold is supplied AND no settings row exists.
DEFAULT_SIM_THRESHOLD = 0.85
_FIGURE_KINDS = ("face", "figure")


async def _settings_threshold(session: AsyncSession) -> float:
    val = (
        await session.execute(
            select(MLSettings.ccip_match_threshold).where(MLSettings.id == 1)
        )
    ).scalar_one_or_none()
    return float(val) if val is not None else DEFAULT_SIM_THRESHOLD


def _l2norm(mat, np):
    n = np.linalg.norm(mat, axis=1, keepdims=True)
    n[n == 0] = 1.0
    return mat / n


async def character_references(session: AsyncSession) -> dict[int, list]:
    """Per character-tag CCIP reference vectors: figure/face-region CCIP
    embeddings on images that carry that character tag (the operator's examples).
    Multi-prototype — several vectors per character."""
    rows = (
        await session.execute(
            select(image_tag.c.tag_id, ImageRegion.ccip_embedding)
            .select_from(ImageRegion)
            .join(
                image_tag,
                image_tag.c.image_record_id == ImageRegion.image_record_id,
            )
            .join(Tag, Tag.id == image_tag.c.tag_id)
            .where(Tag.kind == TagKind.character)
            .where(ImageRegion.kind.in_(_FIGURE_KINDS))
            .where(ImageRegion.ccip_embedding.is_not(None))
        )
    ).all()
    refs: dict[int, list] = {}
    for tag_id, vec in rows:
        refs.setdefault(tag_id, []).append(vec)
    return refs


async def _tag_names(session: AsyncSession, tag_ids: list[int]) -> dict[int, str]:
    if not tag_ids:
        return {}
    return dict(
        (
            await session.execute(
                select(Tag.id, Tag.name).where(Tag.id.in_(tag_ids))
            )
        ).all()
    )


async def match_image(
    session: AsyncSession, image_id: int, threshold: float | None = None
) -> list[dict]:
    """Character suggestions for one image from its figure-region CCIP vectors:
    [{tag_id, name, category:'character', score, source:'ccip'}], ranked.
    Already-applied character tags are excluded. Empty if the image has no figure
    CCIP vectors or no character references exist yet. threshold defaults to the
    live ml_settings.ccip_match_threshold."""
    import numpy as np

    if threshold is None:
        threshold = await _settings_threshold(session)

    qvecs = (
        await session.execute(
            select(ImageRegion.ccip_embedding).where(
                ImageRegion.image_record_id == image_id,
                ImageRegion.kind.in_(_FIGURE_KINDS),
                ImageRegion.ccip_embedding.is_not(None),
            )
        )
    ).scalars().all()
    if not qvecs:
        return []
    refs = await character_references(session)
    if not refs:
        return []
    applied = set(
        (
            await session.execute(
                select(image_tag.c.tag_id).where(
                    image_tag.c.image_record_id == image_id
                )
            )
        ).scalars()
    )
    names = await _tag_names(session, [t for t in refs if t not in applied])

    Q = _l2norm(np.vstack([np.asarray(v, dtype=np.float32) for v in qvecs]), np)
    out = []
    for tag_id, vecs in refs.items():
        if tag_id in applied:
            continue
        R = _l2norm(np.vstack([np.asarray(v, dtype=np.float32) for v in vecs]), np)
        best = float((Q @ R.T).max())   # best (query figure, reference) cosine
        if best >= threshold:
            out.append({
                "tag_id": tag_id,
                "name": names.get(tag_id, str(tag_id)),
                "category": "character",
                "score": round(best, 4),
                "source": "ccip",
            })
    out.sort(key=lambda d: d["score"], reverse=True)
    return out