diff --git a/backend/app/services/ml/suggestions.py b/backend/app/services/ml/suggestions.py index 1cb306f..71fbd24 100644 --- a/backend/app/services/ml/suggestions.py +++ b/backend/app/services/ml/suggestions.py @@ -16,6 +16,7 @@ from sqlalchemy.ext.asyncio import AsyncSession from ...models import ImageRecord, TagSuggestionRejection from ...models.tag import image_tag +from .ccip import match_image as ccip_match_image from .heads import score_image @@ -27,7 +28,7 @@ class Suggestion: display_name: str category: str score: float - source: str # 'head' (Camie 'tagger'/'centroid' sources removed in v2) + source: str # 'head' | 'ccip' | 'both' (Camie tagger/centroid removed in v2) creates_new_tag: bool # raw_name = the booru model vocab key behind this suggestion. It's the key # an alias MUST be stored under (resolution looks up the raw key), so the @@ -92,19 +93,39 @@ class SuggestionService: hits = await score_image( self.session, image_id, threshold_override=threshold_override ) + # CCIP character matches OVERLAY the SigLIP character heads — a + # complementary, identity-specialized signal with different failure modes + # (CCIP needs a detected figure; heads work whole-image). Merged by tag: + # 'both' when they corroborate, taking the higher score. + ccip_hits = await ccip_match_image(self.session, image_id) + + merged: dict[tuple[str, int], dict] = {} + for h in hits: + merged[(h["category"], h["tag_id"])] = { + "name": h["name"], "score": h["score"], "source": "head", + } + for c in ccip_hits: + key = ("character", c["tag_id"]) + ex = merged.get(key) + if ex is not None: + ex["source"] = "both" + ex["score"] = max(ex["score"], c["score"]) + else: + merged[key] = { + "name": c["name"], "score": c["score"], "source": "ccip", + } result = SuggestionList() - for h in hits: - tag_id = h["tag_id"] + for (cat, tag_id), m in merged.items(): if tag_id in applied: continue - result.by_category.setdefault(h["category"], []).append( + result.by_category.setdefault(cat, []).append( Suggestion( canonical_tag_id=tag_id, - display_name=h["name"], - category=h["category"], - score=h["score"], - source="head", + display_name=m["name"], + category=cat, + score=m["score"], + source=m["source"], creates_new_tag=False, rejected=tag_id in rejected, ) diff --git a/tests/test_ml_suggestions.py b/tests/test_ml_suggestions.py index 7a427e2..6d5326e 100644 --- a/tests/test_ml_suggestions.py +++ b/tests/test_ml_suggestions.py @@ -4,7 +4,7 @@ scikit-learn, ml image only); scoring is numpy-only (available via pgvector).""" import pytest from sqlalchemy import select -from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind +from backend.app.models import ImageRecord, ImageRegion, MLSettings, TagHead, TagKind from backend.app.models.tag import image_tag from backend.app.services.ml.allowlist import AllowlistService from backend.app.services.ml.suggestions import SuggestionService @@ -131,3 +131,35 @@ async def test_rejected_tag_surfaced_flagged_then_reversible(db): sl2 = await SuggestionService(db).for_image(img.id) s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id) assert s2.rejected is False + + +async def _figure(db, image_id, slot): + v = [0.0] * 768 + v[slot] = 1.0 + db.add(ImageRegion( + image_record_id=image_id, kind="figure", + rx=0.0, ry=0.0, rw=1.0, rh=1.0, + ccip_embedding=v, embedding_version="ccip-test", + )) + + +@pytest.mark.asyncio +async def test_ccip_character_surfaces_in_rail(db): + # A character with a CCIP reference (a tagged figure) is suggested on a new + # image whose figure matches — overlaid into the rail alongside the heads. + raven = await TagService(db).find_or_create("Raven", TagKind.character) + ref = await _img(db, "0" * 64, None) # the operator's tagged example + await _figure(db, ref.id, slot=0) + await db.execute(image_tag.insert().values( + image_record_id=ref.id, tag_id=raven.id, source="manual", + )) + query = await _img(db, "1" * 64, None) # untagged, matching figure + await _figure(db, query.id, slot=0) + await db.commit() + + sl = await SuggestionService(db).for_image(query.id) + m = next( + c for c in sl.by_category.get("character", []) + if c.canonical_tag_id == raven.id + ) + assert m.source == "ccip"