"""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: # 'artist' intentionally absent (FC-2d-vii-c) — falls through to # the 1.01 "never surfaces" default like any unsurfaced category. return { "character": s.suggestion_threshold_character, "copyright": s.suggestion_threshold_copyright, "general": s.suggestion_threshold_general, }.get(category, 1.01) 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 async def for_selection( self, image_ids: list[int], threshold: float = 0.8, top_k: int = 10, ) -> dict[str, list[dict]]: """Consensus suggestions across image_ids. A tag is included iff it was suggested for (or already applied to) >= threshold fraction of the selection AND was acceptable on >= 1 image. Confidence is the mean over images where it was suggested. Aggregated by canonical_tag_id; creates-new (no canonical id) suggestions are skipped (bulk Accept applies by tag id).""" if not image_ids: return {} threshold = min(1.0, max(0.0, threshold)) total = len(image_ids) stats: dict[int, dict] = {} for image_id in image_ids: sl = await self.for_image(image_id) for category, items in sl.by_category.items(): for s in items: if s.canonical_tag_id is None or s.creates_new_tag: continue st = stats.get(s.canonical_tag_id) if st is None: st = { "tag_id": s.canonical_tag_id, "name": s.display_name, "category": category, "source": s.source, "suggested_count": 0, "sum_score": 0.0, } stats[s.canonical_tag_id] = st st["suggested_count"] += 1 st["sum_score"] += s.score rows = ( await self.session.execute( select( image_tag.c.image_record_id, image_tag.c.tag_id ).where(image_tag.c.image_record_id.in_(image_ids)) ) ).all() applied_by_tag: dict[int, set[int]] = {} for iid, tid in rows: applied_by_tag.setdefault(tid, set()).add(iid) result: dict[str, list[dict]] = {} for st in stats.values(): existing_count = len(applied_by_tag.get(st["tag_id"], set())) covered = st["suggested_count"] + existing_count coverage = covered / total if coverage < threshold or st["suggested_count"] < 1: continue result.setdefault(st["category"], []).append( { "canonical_tag_id": st["tag_id"], "name": st["name"], "category": st["category"], "confidence": round( st["sum_score"] / st["suggested_count"], 4 ), "coverage": round(coverage, 4), "covered_count": covered, "source": st["source"], } ) for cat in result: result[cat].sort(key=lambda x: x["confidence"], reverse=True) result[cat] = result[cat][:top_k] return result