diff --git a/backend/app/services/ml/centroids.py b/backend/app/services/ml/centroids.py new file mode 100644 index 0000000..b4b4f77 --- /dev/null +++ b/backend/app/services/ml/centroids.py @@ -0,0 +1,148 @@ +"""Tag centroids: the mean SigLIP embedding of a tag's member images. + +Powers centroid-augmented suggestions (a tag whose centroid is close to an +image's embedding becomes a suggestion even if Camie didn't predict it). +""" + +from dataclasses import dataclass + +import numpy as np +from sqlalchemy import func, select +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models import ( + ImageRecord, + MLSettings, + Tag, + TagKind, + TagReferenceEmbedding, +) +from ...models.tag import image_tag +from .embedder import MODEL_VERSION as SIGLIP_VERSION + +ELIGIBLE_KINDS = { + TagKind.character, + TagKind.artist, + TagKind.fandom, + TagKind.general, + TagKind.series, +} + + +@dataclass(frozen=True) +class CentroidHit: + tag_id: int + similarity: float + + +class CentroidService: + def __init__(self, session: AsyncSession): + self.session = session + + async def _min_reference_images(self) -> int: + return ( + await self.session.execute( + select(MLSettings.min_reference_images).where(MLSettings.id == 1) + ) + ).scalar_one() + + async def recompute_for_tag(self, tag_id: int) -> bool: + """Recompute one tag's centroid. Returns True if a centroid was + written, False if skipped (ineligible kind or too few members).""" + tag = await self.session.get(Tag, tag_id) + if tag is None or tag.kind not in ELIGIBLE_KINDS: + return False + + min_refs = await self._min_reference_images() + + stmt = ( + select(ImageRecord.siglip_embedding) + .join(image_tag, image_tag.c.image_record_id == ImageRecord.id) + .where(image_tag.c.tag_id == tag_id) + .where(ImageRecord.siglip_embedding.is_not(None)) + ) + embeddings = [ + np.array(e, dtype=np.float32) + for e in (await self.session.execute(stmt)).scalars().all() + ] + if len(embeddings) < min_refs: + return False + + centroid = np.mean(np.stack(embeddings), axis=0).astype(np.float32) + + stmt = insert(TagReferenceEmbedding).values( + tag_id=tag_id, + embedding=centroid.tolist(), + reference_count=len(embeddings), + model_version=SIGLIP_VERSION, + ) + stmt = stmt.on_conflict_do_update( + index_elements=["tag_id"], + set_={ + "embedding": centroid.tolist(), + "reference_count": len(embeddings), + "model_version": SIGLIP_VERSION, + "updated_at": func.now(), + }, + ) + await self.session.execute(stmt) + return True + + async def list_drifted(self) -> list[int]: + """Tag ids whose centroid is stale: member count != reference_count, + OR no centroid row, OR centroid built on a different SigLIP version. + Only considers eligible-kind tags with embeddings present.""" + member_counts = ( + select( + image_tag.c.tag_id.label("tag_id"), + func.count(image_tag.c.image_record_id).label("members"), + ) + .join(ImageRecord, ImageRecord.id == image_tag.c.image_record_id) + .where(ImageRecord.siglip_embedding.is_not(None)) + .group_by(image_tag.c.tag_id) + .subquery() + ) + stmt = ( + select(Tag.id) + .join(member_counts, member_counts.c.tag_id == Tag.id) + .outerjoin( + TagReferenceEmbedding, + TagReferenceEmbedding.tag_id == Tag.id, + ) + .where(Tag.kind.in_(ELIGIBLE_KINDS)) + .where( + (TagReferenceEmbedding.tag_id.is_(None)) + | ( + TagReferenceEmbedding.reference_count + != member_counts.c.members + ) + | (TagReferenceEmbedding.model_version != SIGLIP_VERSION) + ) + ) + return list((await self.session.execute(stmt)).scalars().all()) + + async def find_similar_tags( + self, image_id: int, limit: int = 20 + ) -> list[CentroidHit]: + """Cosine similarity between an image's embedding and stored + centroids. Returns top-`limit` by similarity DESC. pgvector's + cosine_distance gives 1 - cosine_similarity.""" + img = await self.session.get(ImageRecord, image_id) + if img is None or img.siglip_embedding is None: + return [] + emb = img.siglip_embedding + distance = TagReferenceEmbedding.embedding.cosine_distance(emb) + stmt = ( + select( + TagReferenceEmbedding.tag_id, + (1 - distance).label("similarity"), + ) + .order_by(distance.asc()) + .limit(limit) + ) + rows = (await self.session.execute(stmt)).all() + return [ + CentroidHit(tag_id=r.tag_id, similarity=float(r.similarity)) + for r in rows + ] diff --git a/tests/test_ml_centroids.py b/tests/test_ml_centroids.py new file mode 100644 index 0000000..192bd1c --- /dev/null +++ b/tests/test_ml_centroids.py @@ -0,0 +1,112 @@ +import numpy as np +import pytest + +from backend.app.models import ImageRecord, TagKind +from backend.app.models.tag import image_tag +from backend.app.services.ml.centroids import CentroidService +from backend.app.services.tag_service import TagService + +pytestmark = pytest.mark.integration + + +def _img(sha: str, embedding: list[float] | None) -> ImageRecord: + return ImageRecord( + path=f"/images/{sha}.jpg", + sha256=sha, + size_bytes=1, + mime="image/jpeg", + width=1, + height=1, + origin="imported_filesystem", + integrity_status="unknown", + siglip_embedding=embedding, + ) + + +async def _attach(db, image_id: int, tag_id: int): + await db.execute( + image_tag.insert().values( + image_record_id=image_id, tag_id=tag_id, source="manual" + ) + ) + + +@pytest.mark.asyncio +async def test_recompute_skips_too_few_members(db): + tags = TagService(db) + tag = await tags.find_or_create("Lonely", TagKind.character) + img = _img("a" * 64, [0.1] * 1152) + db.add(img) + await db.flush() + await _attach(db, img.id, tag.id) + + svc = CentroidService(db) + assert await svc.recompute_for_tag(tag.id) is False + + +@pytest.mark.asyncio +async def test_recompute_writes_centroid(db): + tags = TagService(db) + tag = await tags.find_or_create("Popular", TagKind.character) + for i in range(5): + img = _img(f"{i:064d}", [float(i)] * 1152) + db.add(img) + await db.flush() + await _attach(db, img.id, tag.id) + + svc = CentroidService(db) + assert await svc.recompute_for_tag(tag.id) is True + + from backend.app.models import TagReferenceEmbedding + cen = await db.get(TagReferenceEmbedding, tag.id) + assert cen is not None + assert cen.reference_count == 5 + assert abs(np.array(cen.embedding)[0] - 2.0) < 1e-4 + + +@pytest.mark.asyncio +async def test_recompute_skips_ineligible_kind(db): + tags = TagService(db) + tag = await tags.find_or_create("somearchive", TagKind.archive) + for i in range(5): + img = _img(f"arch{i:060d}", [1.0] * 1152) + db.add(img) + await db.flush() + await _attach(db, img.id, tag.id) + svc = CentroidService(db) + assert await svc.recompute_for_tag(tag.id) is False + + +@pytest.mark.asyncio +async def test_list_drifted_includes_uncomputed(db): + tags = TagService(db) + tag = await tags.find_or_create("Drifty", TagKind.character) + for i in range(5): + img = _img(f"d{i:063d}", [0.5] * 1152) + db.add(img) + await db.flush() + await _attach(db, img.id, tag.id) + svc = CentroidService(db) + drifted = await svc.list_drifted() + assert tag.id in drifted + + +@pytest.mark.asyncio +async def test_find_similar_tags(db): + tags = TagService(db) + tag = await tags.find_or_create("SimTag", TagKind.character) + for i in range(5): + img = _img(f"s{i:063d}", [1.0] * 1152) + db.add(img) + await db.flush() + await _attach(db, img.id, tag.id) + svc = CentroidService(db) + await svc.recompute_for_tag(tag.id) + + query_img = _img("q" * 64, [1.0] * 1152) + db.add(query_img) + await db.flush() + hits = await svc.find_similar_tags(query_img.id, limit=10) + assert any(h.tag_id == tag.id for h in hits) + sim = next(h.similarity for h in hits if h.tag_id == tag.id) + assert sim > 0.99