bvandeusen
dfe2fda564
match_image now tracks WHICH query figure produced the winning cosine per
character (argmax over the per-figure best-reference sim) and attaches its bbox as
grounding {bbox,kind:'figure',detector}. SuggestionService carries it: a CCIP-only
character hit grounds to its figure; a 'both' hit keeps the head's localized crop
if it had one, else falls back to the CCIP figure — so corroborated characters
stay grounded. Test: a character match carries the matched figure's bbox+kind.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01CDgx8bQS5YrGRK76v8HUnM
228 lines
9.6 KiB
Python
228 lines
9.6 KiB
Python
"""The suggestion read-path: trained HEADS score one image's frozen embedding
|
|
into alias-resolved, category-grouped, ranked suggestions.
|
|
|
|
Tagging-v2 (#114): suggestions now come from the per-concept heads that LEARN
|
|
from the operator's tags (services/ml/heads.py) — the Camie prediction source
|
|
and the per-tag SigLIP centroid have been REMOVED. A head exists only for an
|
|
existing concept tag, so every suggestion is a canonical tag (no raw model key,
|
|
no alias remap, no creates-new). Rejected tags stay in the list FLAGGED (not
|
|
dropped) so the rail can show + reverse a dismissal.
|
|
"""
|
|
|
|
from dataclasses import dataclass, field
|
|
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from ...models import ImageRecord, TagSuggestionRejection
|
|
from ...models.tag import image_tag
|
|
from .ccip import match_image as ccip_match_image
|
|
from .heads import score_image
|
|
|
|
|
|
@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 # 'head' | 'ccip' | 'both' (Camie tagger/centroid removed in v2)
|
|
creates_new_tag: bool
|
|
# raw_name = the booru model vocab key behind this suggestion. It's the key
|
|
# an alias MUST be stored under (resolution looks up the raw key), so the
|
|
# modal needs it to author an alias correctly. None for centroid-only hits
|
|
# (no underlying prediction → nothing to alias).
|
|
raw_name: str | None = None
|
|
# via_alias = this suggestion was surfaced because an operator alias remapped
|
|
# the raw prediction to this canonical tag. Lets the UI mark it + offer undo.
|
|
via_alias: bool = False
|
|
# rejected = the operator dismissed this tag for this image (a stored
|
|
# TagSuggestionRejection). It stays in the list — flagged, not dropped — so
|
|
# the rejection is VISIBLE and REVERSIBLE in the rail (misclick recovery,
|
|
# operator-asked 2026-06-27) instead of silently vanishing or re-suggesting.
|
|
rejected: bool = False
|
|
# grounding = the crop region that produced this suggestion (#1206):
|
|
# {bbox:[x,y,w,h] normalized, kind, detector}. None when the whole-image
|
|
# vector won (not localized) or for a CCIP-only hit (figure grounding TBD).
|
|
# Lets the rail highlight the exact region on hover.
|
|
grounding: dict | None = None
|
|
|
|
|
|
@dataclass
|
|
class SuggestionList:
|
|
by_category: dict[str, list[Suggestion]] = field(default_factory=dict)
|
|
|
|
|
|
class SuggestionService:
|
|
def __init__(self, session: AsyncSession):
|
|
self.session = session
|
|
|
|
async def for_image(
|
|
self, image_id: int, threshold_override: float | None = None,
|
|
) -> SuggestionList:
|
|
"""Head-scored suggestions for one image, grouped by category and ranked.
|
|
|
|
Each trained head scores the image's frozen embedding; a concept surfaces
|
|
when its score clears the head's own suggest threshold. threshold_override
|
|
(used by the typed tag-input dropdown's "show everything" mode) replaces
|
|
that per-head cut with a flat floor (0 → every head), so a low-scoring
|
|
concept can still be typed + picked in canonical formatting.
|
|
|
|
Already-applied tags are dropped; rejected tags stay FLAGGED and sink to
|
|
the bottom of their category so a dismissal is visible + reversible."""
|
|
img = await self.session.get(ImageRecord, image_id)
|
|
if img is None:
|
|
return SuggestionList()
|
|
|
|
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()
|
|
)
|
|
|
|
hits = await score_image(
|
|
self.session, image_id, threshold_override=threshold_override
|
|
)
|
|
# CCIP character matches OVERLAY the SigLIP character heads — a
|
|
# complementary, identity-specialized signal with different failure modes
|
|
# (CCIP needs a detected figure; heads work whole-image). Merged by tag:
|
|
# 'both' when they corroborate, taking the higher score.
|
|
ccip_hits = await ccip_match_image(self.session, image_id)
|
|
|
|
merged: dict[tuple[str, int], dict] = {}
|
|
for h in hits:
|
|
merged[(h["category"], h["tag_id"])] = {
|
|
"name": h["name"], "score": h["score"], "source": "head",
|
|
"grounding": h.get("grounding"),
|
|
}
|
|
for c in ccip_hits:
|
|
key = ("character", c["tag_id"])
|
|
ex = merged.get(key)
|
|
if ex is not None:
|
|
ex["source"] = "both"
|
|
ex["score"] = max(ex["score"], c["score"])
|
|
# Keep the head's localized crop if it had one; else fall back to
|
|
# the CCIP figure so a corroborated character still grounds (#1206).
|
|
ex["grounding"] = ex.get("grounding") or c.get("grounding")
|
|
else:
|
|
merged[key] = {
|
|
"name": c["name"], "score": c["score"], "source": "ccip",
|
|
"grounding": c.get("grounding"),
|
|
}
|
|
|
|
result = SuggestionList()
|
|
for (cat, tag_id), m in merged.items():
|
|
if tag_id in applied:
|
|
continue
|
|
result.by_category.setdefault(cat, []).append(
|
|
Suggestion(
|
|
canonical_tag_id=tag_id,
|
|
display_name=m["name"],
|
|
category=cat,
|
|
score=m["score"],
|
|
source=m["source"],
|
|
creates_new_tag=False,
|
|
rejected=tag_id in rejected,
|
|
grounding=m.get("grounding"),
|
|
)
|
|
)
|
|
for cat in result.by_category:
|
|
# Live suggestions first (by score), rejected ones sink to the
|
|
# bottom of the category — visible for recovery, out of the way.
|
|
result.by_category[cat].sort(key=lambda s: (s.rejected, -s.score))
|
|
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
|
|
# for_image keeps rejected tags (flagged) for the rail;
|
|
# bulk consensus must still ignore them — a tag dismissed on
|
|
# an image isn't a suggestion for that image.
|
|
if s.rejected:
|
|
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
|