bvandeusen
22cdf0f334
Switch every prediction READER off the JSON column onto the normalized
image_prediction table. Parity by construction: each reader loads the same
{raw_name: {category, confidence}} dict it consumed before (via small
_load_predictions helpers), so all downstream threshold/alias/merge/consensus
logic is byte-identical — only the data source changed.
- suggestions.SuggestionService.for_image (and for_selection via it)
- ml.apply_allowlist_tags (iterates images that have prediction rows)
- importer re-import reset deletes the image's prediction rows
The tagger_predictions JSON column is still dual-written (step 1) so it stays
valid during transition; the backfill task's NULL check still works. Removing
the JSON write + DROP column + retiring the #764 prune is the cleanup
follow-up (needs a quiesced-worker window for the DROP lock).
Tests: shared tests/_prediction_helpers.seed_predictions seeds the table;
read-path tests (suggestions, bulk consensus, allowlist apply, API) seed there
instead of ImageRecord.tagger_predictions.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
330 lines
13 KiB
Python
330 lines
13 KiB
Python
"""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 func, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from ...models import (
|
|
ImagePrediction,
|
|
ImageRecord,
|
|
MLSettings,
|
|
Tag,
|
|
TagSuggestionRejection,
|
|
)
|
|
from ...models.tag import image_tag
|
|
from .aliases import AliasService
|
|
from .centroids import CentroidService
|
|
from .tag_name import normalize as normalize_tag_name
|
|
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()
|
|
|
|
async def _load_predictions(self, image_id: int) -> dict:
|
|
"""Predictions for one image from the normalized image_prediction
|
|
table (#768), in the {raw_name: {category, confidence}} shape the rest
|
|
of this service consumed from the old JSON column — so all downstream
|
|
threshold/alias/merge logic is unchanged."""
|
|
rows = (
|
|
await self.session.execute(
|
|
select(
|
|
ImagePrediction.raw_name,
|
|
ImagePrediction.category,
|
|
ImagePrediction.score,
|
|
).where(ImagePrediction.image_record_id == image_id)
|
|
)
|
|
).all()
|
|
return {
|
|
r.raw_name: {"category": r.category, "confidence": r.score}
|
|
for r in rows
|
|
}
|
|
|
|
def _threshold_for(
|
|
self, s: MLSettings, category: str, override: float | None = None,
|
|
) -> float:
|
|
# 'artist' (FC-2d-vii-c) and 'copyright' (2026-06-01) retired;
|
|
# both fall through to the 1.01 "never surfaces" default like any
|
|
# unsurfaced category.
|
|
# override (the typed-dropdown "show everything the model saw" mode)
|
|
# applies to the surfaced categories only — unsurfaced ones are already
|
|
# skipped before the threshold check, so they can't leak in.
|
|
if override is not None:
|
|
return override
|
|
return {
|
|
"character": s.suggestion_threshold_character,
|
|
"general": s.suggestion_threshold_general,
|
|
}.get(category, 1.01)
|
|
|
|
async def for_image(
|
|
self, image_id: int, *, threshold_override: float | None = None,
|
|
) -> SuggestionList:
|
|
"""Ranked suggestions for one image.
|
|
|
|
threshold_override surfaces EVERY stored tagger prediction (down to the
|
|
ingest STORE_FLOOR) regardless of the configured per-category suggestion
|
|
thresholds — backs the tag-input dropdown's "search all of the model's
|
|
predictions, including low-confidence ones, in the canonical formatting"
|
|
mode (operator-asked 2026-06-09). The Suggestions panel still calls with
|
|
no override so it stays the curated above-threshold list."""
|
|
img = await self.session.get(ImageRecord, image_id)
|
|
if img is None:
|
|
return SuggestionList()
|
|
|
|
settings = await self._settings()
|
|
predictions: dict = await self._load_predictions(image_id)
|
|
|
|
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 carry (raw_name, display_name, category, confidence).
|
|
# raw_name = the booru-formatted vocab key, kept for alias_map
|
|
# lookup since alias rows are hand-curated against raw keys.
|
|
# display_name = normalize_tag_name(raw_name) — what the operator
|
|
# sees AND what gets written to tag.name on Accept.
|
|
candidates: list[tuple[str, 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, threshold_override):
|
|
continue
|
|
display = normalize_tag_name(name)
|
|
if display is None:
|
|
# emoticon / pure-punctuation vocab entry — drop entirely
|
|
continue
|
|
candidates.append((name, display, category, conf))
|
|
|
|
alias_map = await self.aliases.resolve_many(
|
|
[(raw, c) for raw, _disp, c, _conf 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 raw, display, category, conf in candidates:
|
|
canonical = alias_map.get((raw, 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:
|
|
# Case-insensitive match on BOTH the raw camie key AND
|
|
# the normalized form — covers legacy underscore-named
|
|
# Tag rows accepted before normalization shipped, AND
|
|
# any tag the operator created with the human form.
|
|
existing_tag = (
|
|
await self.session.execute(
|
|
select(Tag).where(
|
|
func.lower(Tag.name).in_(
|
|
[raw.lower(), display.lower()]
|
|
)
|
|
)
|
|
)
|
|
).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:{display}:{category}",
|
|
Suggestion(
|
|
canonical_tag_id=None,
|
|
display_name=display,
|
|
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
|