feat(fc2b): add SuggestionService — alias-resolved, threshold-filtered, ranked
The read path: load tagger_predictions, drop unsurfaced categories (rating/meta/year), apply per-category thresholds, batch-resolve aliases, skip applied + rejected, augment with centroid hits above the similarity threshold, merge duplicate signals (take max score, mark source 'both'), group by category, sort by score DESC. Tests marked integration. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,199 @@
|
||||
"""The suggestion read-path: raw predictions + centroids -> alias-resolved,
|
||||
threshold-filtered, category-grouped, ranked suggestions for one image.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from ...models import (
|
||||
ImageRecord,
|
||||
MLSettings,
|
||||
Tag,
|
||||
TagSuggestionRejection,
|
||||
)
|
||||
from ...models.tag import image_tag
|
||||
from .aliases import AliasService
|
||||
from .centroids import CentroidService
|
||||
from .tagger import SURFACED_CATEGORIES
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Suggestion:
|
||||
# canonical_tag_id is None when this is a raw Camie tag with no alias and
|
||||
# no existing Tag row — accepting it will create the tag.
|
||||
canonical_tag_id: int | None
|
||||
display_name: str
|
||||
category: str
|
||||
score: float
|
||||
source: str # 'tagger' | 'centroid' | 'both'
|
||||
creates_new_tag: bool
|
||||
|
||||
|
||||
@dataclass
|
||||
class SuggestionList:
|
||||
by_category: dict[str, list[Suggestion]] = field(default_factory=dict)
|
||||
|
||||
|
||||
class SuggestionService:
|
||||
def __init__(self, session: AsyncSession):
|
||||
self.session = session
|
||||
self.aliases = AliasService(session)
|
||||
self.centroids = CentroidService(session)
|
||||
|
||||
async def _settings(self) -> MLSettings:
|
||||
return (
|
||||
await self.session.execute(select(MLSettings).where(MLSettings.id == 1))
|
||||
).scalar_one()
|
||||
|
||||
def _threshold_for(self, s: MLSettings, category: str) -> float:
|
||||
return {
|
||||
"artist": s.suggestion_threshold_artist,
|
||||
"character": s.suggestion_threshold_character,
|
||||
"copyright": s.suggestion_threshold_copyright,
|
||||
"general": s.suggestion_threshold_general,
|
||||
}.get(category, 1.01) # 1.01 => never surfaces (unsurfaced category)
|
||||
|
||||
async def for_image(self, image_id: int) -> SuggestionList:
|
||||
img = await self.session.get(ImageRecord, image_id)
|
||||
if img is None:
|
||||
return SuggestionList()
|
||||
|
||||
settings = await self._settings()
|
||||
predictions: dict = img.tagger_predictions or {}
|
||||
|
||||
applied = set(
|
||||
(
|
||||
await self.session.execute(
|
||||
select(image_tag.c.tag_id).where(
|
||||
image_tag.c.image_record_id == image_id
|
||||
)
|
||||
)
|
||||
).scalars().all()
|
||||
)
|
||||
rejected = set(
|
||||
(
|
||||
await self.session.execute(
|
||||
select(TagSuggestionRejection.tag_id).where(
|
||||
TagSuggestionRejection.image_record_id == image_id
|
||||
)
|
||||
)
|
||||
).scalars().all()
|
||||
)
|
||||
|
||||
# --- Camie predictions ---
|
||||
candidates: list[tuple[str, str, float]] = []
|
||||
for name, p in predictions.items():
|
||||
category = p.get("category", "general")
|
||||
if category not in SURFACED_CATEGORIES:
|
||||
continue
|
||||
conf = float(p.get("confidence", 0.0))
|
||||
if conf < self._threshold_for(settings, category):
|
||||
continue
|
||||
candidates.append((name, category, conf))
|
||||
|
||||
alias_map = await self.aliases.resolve_many(
|
||||
[(n, c) for n, c, _ in candidates]
|
||||
)
|
||||
|
||||
merged: dict[object, Suggestion] = {}
|
||||
|
||||
def _merge(key, sug: Suggestion):
|
||||
existing = merged.get(key)
|
||||
if existing is None:
|
||||
merged[key] = sug
|
||||
elif sug.score > existing.score:
|
||||
merged[key] = Suggestion(
|
||||
canonical_tag_id=existing.canonical_tag_id,
|
||||
display_name=existing.display_name,
|
||||
category=existing.category,
|
||||
score=sug.score,
|
||||
source="both"
|
||||
if existing.source != sug.source
|
||||
else existing.source,
|
||||
creates_new_tag=existing.creates_new_tag,
|
||||
)
|
||||
|
||||
for name, category, conf in candidates:
|
||||
canonical = alias_map.get((name, category))
|
||||
if canonical is not None:
|
||||
if canonical.id in applied or canonical.id in rejected:
|
||||
continue
|
||||
_merge(
|
||||
canonical.id,
|
||||
Suggestion(
|
||||
canonical_tag_id=canonical.id,
|
||||
display_name=canonical.name,
|
||||
category=category,
|
||||
score=conf,
|
||||
source="tagger",
|
||||
creates_new_tag=False,
|
||||
),
|
||||
)
|
||||
else:
|
||||
existing_tag = (
|
||||
await self.session.execute(
|
||||
select(Tag).where(Tag.name == name)
|
||||
)
|
||||
).scalars().first()
|
||||
if existing_tag is not None:
|
||||
if (
|
||||
existing_tag.id in applied
|
||||
or existing_tag.id in rejected
|
||||
):
|
||||
continue
|
||||
_merge(
|
||||
existing_tag.id,
|
||||
Suggestion(
|
||||
canonical_tag_id=existing_tag.id,
|
||||
display_name=existing_tag.name,
|
||||
category=category,
|
||||
score=conf,
|
||||
source="tagger",
|
||||
creates_new_tag=False,
|
||||
),
|
||||
)
|
||||
else:
|
||||
_merge(
|
||||
f"raw:{name}:{category}",
|
||||
Suggestion(
|
||||
canonical_tag_id=None,
|
||||
display_name=name,
|
||||
category=category,
|
||||
score=conf,
|
||||
source="tagger",
|
||||
creates_new_tag=True,
|
||||
),
|
||||
)
|
||||
|
||||
# --- Centroid augmentation ---
|
||||
hits = await self.centroids.find_similar_tags(image_id, limit=30)
|
||||
for hit in hits:
|
||||
if hit.similarity < settings.centroid_similarity_threshold:
|
||||
continue
|
||||
if hit.tag_id in applied or hit.tag_id in rejected:
|
||||
continue
|
||||
tag = await self.session.get(Tag, hit.tag_id)
|
||||
if tag is None:
|
||||
continue
|
||||
cat = tag.kind.value if hasattr(tag.kind, "value") else str(tag.kind)
|
||||
display_cat = cat if cat in SURFACED_CATEGORIES else "general"
|
||||
_merge(
|
||||
tag.id,
|
||||
Suggestion(
|
||||
canonical_tag_id=tag.id,
|
||||
display_name=tag.name,
|
||||
category=display_cat,
|
||||
score=hit.similarity,
|
||||
source="centroid",
|
||||
creates_new_tag=False,
|
||||
),
|
||||
)
|
||||
|
||||
result = SuggestionList()
|
||||
for sug in merged.values():
|
||||
result.by_category.setdefault(sug.category, []).append(sug)
|
||||
for cat in result.by_category:
|
||||
result.by_category[cat].sort(key=lambda s: s.score, reverse=True)
|
||||
return result
|
||||
@@ -0,0 +1,105 @@
|
||||
import pytest
|
||||
|
||||
from backend.app.models import ImageRecord, TagKind
|
||||
from backend.app.models.tag import image_tag
|
||||
from backend.app.services.ml.aliases import AliasService
|
||||
from backend.app.services.ml.suggestions import SuggestionService
|
||||
from backend.app.services.tag_service import TagService
|
||||
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
def _img(sha: str, predictions: dict) -> 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",
|
||||
tagger_predictions=predictions,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_threshold_filters_low_confidence_general(db):
|
||||
img = _img(
|
||||
"a" * 64,
|
||||
{
|
||||
"smile": {"category": "general", "confidence": 0.80},
|
||||
"sword": {"category": "general", "confidence": 0.97},
|
||||
},
|
||||
)
|
||||
db.add(img)
|
||||
await db.flush()
|
||||
sl = await SuggestionService(db).for_image(img.id)
|
||||
names = [s.display_name for s in sl.by_category.get("general", [])]
|
||||
assert "sword" in names
|
||||
assert "smile" not in names
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unsurfaced_category_dropped(db):
|
||||
img = _img(
|
||||
"b" * 64,
|
||||
{"safe": {"category": "rating", "confidence": 0.99}},
|
||||
)
|
||||
db.add(img)
|
||||
await db.flush()
|
||||
sl = await SuggestionService(db).for_image(img.id)
|
||||
assert "rating" not in sl.by_category
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_alias_resolution(db):
|
||||
tags = TagService(db)
|
||||
canonical = await tags.find_or_create("Sasuke Uchiha", TagKind.character)
|
||||
await AliasService(db).create("uchiha_sasuke", "character", canonical.id)
|
||||
img = _img(
|
||||
"c" * 64,
|
||||
{"uchiha_sasuke": {"category": "character", "confidence": 0.96}},
|
||||
)
|
||||
db.add(img)
|
||||
await db.flush()
|
||||
sl = await SuggestionService(db).for_image(img.id)
|
||||
chars = sl.by_category["character"]
|
||||
assert len(chars) == 1
|
||||
assert chars[0].display_name == "Sasuke Uchiha"
|
||||
assert chars[0].canonical_tag_id == canonical.id
|
||||
assert chars[0].creates_new_tag is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_tag_creates_new(db):
|
||||
img = _img(
|
||||
"d" * 64,
|
||||
{"brand_new_tag": {"category": "character", "confidence": 0.96}},
|
||||
)
|
||||
db.add(img)
|
||||
await db.flush()
|
||||
sl = await SuggestionService(db).for_image(img.id)
|
||||
chars = sl.by_category["character"]
|
||||
assert chars[0].display_name == "brand_new_tag"
|
||||
assert chars[0].creates_new_tag is True
|
||||
assert chars[0].canonical_tag_id is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_applied_tag_not_suggested(db):
|
||||
tags = TagService(db)
|
||||
tag = await tags.find_or_create("alreadyhere", TagKind.character)
|
||||
img = _img(
|
||||
"e" * 64,
|
||||
{"alreadyhere": {"category": "character", "confidence": 0.96}},
|
||||
)
|
||||
db.add(img)
|
||||
await db.flush()
|
||||
await db.execute(
|
||||
image_tag.insert().values(
|
||||
image_record_id=img.id, tag_id=tag.id, source="manual"
|
||||
)
|
||||
)
|
||||
sl = await SuggestionService(db).for_image(img.id)
|
||||
assert "character" not in sl.by_category or not sl.by_category["character"]
|
||||
Reference in New Issue
Block a user