feat(fc2b): add CentroidService — per-tag SigLIP centroids + similarity
recompute_for_tag (mean of member embeddings, eligible-kind + min-refs gated, upsert), list_drifted (the delta-gate: member-count mismatch OR missing OR wrong model version), find_similar_tags (pgvector cosine distance, similarity = 1 - distance). Tests marked integration. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user