feat: ML tag suggestions, character/fandom integrity, underscores, modal polish
Consolidated merge of feat/tag-suggestions branch. Original 64-commit history was lost to git-object corruption in a Nextcloud-synced checkout; this single commit captures the equivalent diff. Includes: - pgvector-backed tag suggestion infra (WD14 + SigLIP centroids, ml-worker container, Celery tasks, suggestion service, accept/reject endpoints + modal UI with green/red chip buttons) - Character/fandom integrity: title-case normalization on every write path, fandom-id backfill, maintenance task + settings button, migrations g26041901 + h26041901 to canonicalize legacy rows with case-only duplicate merging - Tag-underscores + modal polish: WD14 name canonicalization at emit + accept + add/bulk-add paths, migration i26041901 for legacy-row rename-or-merge across character/fandom/NULL kinds, suggestion-accept refresh parity via awaited loadTags, persistent chip tint
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""Read-path service modules called from Flask routes."""
|
||||
@@ -0,0 +1,235 @@
|
||||
"""Compute merged tag suggestions for an image on demand.
|
||||
|
||||
Sources:
|
||||
- WD14 predictions filtered by per-category thresholds
|
||||
- Embedding similarity vs. character centroids (for character tags only)
|
||||
|
||||
The result is grouped by category and annotated with whether each tag already
|
||||
exists in the Tag table (the modal dims disallowed auto-creations).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
from typing import Iterable
|
||||
|
||||
from sqlalchemy import select
|
||||
|
||||
from app import db
|
||||
from app.models import (
|
||||
ImageTagPrediction,
|
||||
ImageEmbedding,
|
||||
TagReferenceEmbedding,
|
||||
TagSuggestionConfig,
|
||||
ImageRecord,
|
||||
Tag,
|
||||
image_tags,
|
||||
)
|
||||
|
||||
# Tag kinds that receive embedding-similarity suggestions. None = general/topic tags.
|
||||
# Adding 'artist' or 'series' here enables them with no other code changes.
|
||||
ELIGIBLE_CENTROID_KINDS = ('character', 'fandom', None)
|
||||
|
||||
# Maps the WD14 prediction category ('character', 'copyright') to the kind used
|
||||
# in the Tag table. 'general' maps to None and is handled separately at accept
|
||||
# time. Lives here (not in main) so the canonicalization helper below can use it
|
||||
# without a circular import.
|
||||
_WD14_CATEGORY_TO_KIND = {
|
||||
'character': 'character',
|
||||
'copyright': 'fandom',
|
||||
}
|
||||
|
||||
|
||||
def _canonicalize_wd14_name(raw: str, category: str) -> str:
|
||||
"""Rewrite a raw WD14 tag name to the canonical form ImageRepo persists.
|
||||
|
||||
- 'character' / 'copyright': title-case the display portion.
|
||||
- everything else ('general', 'meta', unknown): underscore-to-space only,
|
||||
preserving the user's casing for general topics.
|
||||
"""
|
||||
from app.utils.tag_names import normalize_display_name, underscores_to_spaces
|
||||
kind = _WD14_CATEGORY_TO_KIND.get(category)
|
||||
if ':' in raw:
|
||||
prefix, rest = raw.split(':', 1)
|
||||
transform = normalize_display_name if kind in ('character', 'fandom') else underscores_to_spaces
|
||||
return f"{prefix}:{transform(rest)}"
|
||||
return normalize_display_name(raw) if kind in ('character', 'fandom') else underscores_to_spaces(raw)
|
||||
|
||||
|
||||
# Config keys with safe defaults (used if the row is missing from tag_suggestion_config).
|
||||
_DEFAULTS = {
|
||||
'threshold_general': 0.35,
|
||||
'threshold_character_wd14': 0.75,
|
||||
'threshold_copyright': 0.5,
|
||||
'threshold_meta': 0.5,
|
||||
'threshold_embedding': 0.85,
|
||||
'min_reference_images': 5.0,
|
||||
'wd14_model_version': 'wd-eva02-large-tagger-v3',
|
||||
'siglip_model_version': 'siglip-so400m-patch14-384',
|
||||
}
|
||||
|
||||
|
||||
def _config() -> dict:
|
||||
rows = TagSuggestionConfig.query.all()
|
||||
cfg = dict(_DEFAULTS)
|
||||
for r in rows:
|
||||
try:
|
||||
cfg[r.key] = float(r.value)
|
||||
except ValueError:
|
||||
cfg[r.key] = r.value
|
||||
return cfg
|
||||
|
||||
|
||||
def _category_threshold(category: str, cfg: dict) -> float:
|
||||
return {
|
||||
'general': cfg['threshold_general'],
|
||||
'character': cfg['threshold_character_wd14'],
|
||||
'copyright': cfg['threshold_copyright'],
|
||||
'meta': cfg['threshold_meta'],
|
||||
}.get(category, 1.01) # unknown/rating → effectively disabled (not surfaced)
|
||||
|
||||
|
||||
def _existing_tag_names_for_image(image_id: int) -> set[str]:
|
||||
"""Names of tags already attached to this image.
|
||||
|
||||
Names are returned as-stored; since every write path canonicalizes, any
|
||||
row in the DB is already in canonical form. WD14 output is canonicalized
|
||||
before lookup, so equality works.
|
||||
"""
|
||||
rows = (
|
||||
db.session.query(Tag.name)
|
||||
.join(image_tags, image_tags.c.tag_id == Tag.id)
|
||||
.filter(image_tags.c.image_id == image_id)
|
||||
.all()
|
||||
)
|
||||
return {row[0] for row in rows}
|
||||
|
||||
|
||||
def _existing_tag_names() -> set[str]:
|
||||
return {row[0] for row in db.session.query(Tag.name).all()}
|
||||
|
||||
|
||||
def _wd14_suggestions(image_id: int, cfg: dict, already: set[str]) -> list[dict]:
|
||||
wd14_ver = cfg.get('wd14_model_version', _DEFAULTS['wd14_model_version'])
|
||||
preds = (
|
||||
ImageTagPrediction.query
|
||||
.filter_by(image_id=image_id, model_version=wd14_ver)
|
||||
.all()
|
||||
)
|
||||
out: list[dict] = []
|
||||
for p in preds:
|
||||
threshold = _category_threshold(p.tag_category, cfg)
|
||||
if p.confidence < threshold:
|
||||
continue
|
||||
canonical = _canonicalize_wd14_name(p.tag_name, p.tag_category)
|
||||
if canonical in already:
|
||||
continue
|
||||
out.append({
|
||||
'name': canonical,
|
||||
'category': p.tag_category,
|
||||
'confidence': p.confidence,
|
||||
'source': 'wd14',
|
||||
})
|
||||
return out
|
||||
|
||||
|
||||
# Maps tag kind (as stored on Tag and TagReferenceEmbedding) to the display
|
||||
# category used in the modal UI's grouped suggestions.
|
||||
_KIND_TO_DISPLAY_CATEGORY = {
|
||||
'character': 'character',
|
||||
'fandom': 'copyright',
|
||||
None: 'general',
|
||||
}
|
||||
|
||||
|
||||
def _embedding_tag_suggestions(image_id: int, cfg: dict, already: set[str]) -> list[dict]:
|
||||
siglip_ver = cfg.get('siglip_model_version', _DEFAULTS['siglip_model_version'])
|
||||
min_refs = int(cfg.get('min_reference_images', 5))
|
||||
threshold = cfg.get('threshold_embedding', 0.85)
|
||||
|
||||
image_emb = (
|
||||
ImageEmbedding.query
|
||||
.filter_by(image_id=image_id, model_version=siglip_ver)
|
||||
.first()
|
||||
)
|
||||
if image_emb is None:
|
||||
return []
|
||||
|
||||
# pgvector cosine distance; similarity = 1 - distance
|
||||
distance = TagReferenceEmbedding.centroid.cosine_distance(image_emb.embedding)
|
||||
rows = (
|
||||
db.session.query(
|
||||
TagReferenceEmbedding.tag_name,
|
||||
TagReferenceEmbedding.tag_kind,
|
||||
TagReferenceEmbedding.reference_count,
|
||||
distance.label('distance'),
|
||||
)
|
||||
.filter(TagReferenceEmbedding.model_version == siglip_ver)
|
||||
.filter(TagReferenceEmbedding.reference_count >= min_refs)
|
||||
.order_by(distance)
|
||||
.limit(40)
|
||||
.all()
|
||||
)
|
||||
out: list[dict] = []
|
||||
for tag_name, tag_kind, ref_count, dist in rows:
|
||||
similarity = 1.0 - float(dist)
|
||||
if similarity < threshold:
|
||||
continue
|
||||
if tag_name in already:
|
||||
continue
|
||||
category = _KIND_TO_DISPLAY_CATEGORY.get(tag_kind)
|
||||
if category is None and tag_kind is not None:
|
||||
# Unknown kind — skip defensively. (Should never happen because
|
||||
# recompute_centroid only writes eligible kinds.)
|
||||
continue
|
||||
out.append({
|
||||
'name': tag_name,
|
||||
'category': category,
|
||||
'confidence': similarity,
|
||||
'source': 'embedding_similarity',
|
||||
})
|
||||
return out
|
||||
|
||||
|
||||
def _merge(wd14: list[dict], embedding: list[dict]) -> list[dict]:
|
||||
by_name: dict[str, dict] = {}
|
||||
for s in wd14 + embedding:
|
||||
existing = by_name.get(s['name'])
|
||||
if existing is None or s['confidence'] > existing['confidence']:
|
||||
by_name[s['name']] = s
|
||||
return list(by_name.values())
|
||||
|
||||
|
||||
def get_suggestions(image_id: int, top_k_per_category: int = 10) -> dict:
|
||||
"""Return suggestions grouped by category for one image.
|
||||
|
||||
Shape:
|
||||
{
|
||||
'character': [{name, confidence, source, exists_in_db}, ...],
|
||||
'copyright': [...],
|
||||
'general': [...],
|
||||
}
|
||||
Ordered by confidence desc within each category. 'meta' and 'rating' are omitted.
|
||||
"""
|
||||
if ImageRecord.query.get(image_id) is None:
|
||||
return {'character': [], 'copyright': [], 'general': []}
|
||||
|
||||
cfg = _config()
|
||||
already = _existing_tag_names_for_image(image_id)
|
||||
|
||||
merged = _merge(
|
||||
_wd14_suggestions(image_id, cfg, already),
|
||||
_embedding_tag_suggestions(image_id, cfg, already),
|
||||
)
|
||||
|
||||
existing_all = _existing_tag_names()
|
||||
|
||||
grouped: dict[str, list[dict]] = {'character': [], 'copyright': [], 'general': []}
|
||||
for s in merged:
|
||||
if s['category'] not in grouped:
|
||||
continue # meta / rating / artist / unknown are dropped
|
||||
s['exists_in_db'] = s['name'] in existing_all
|
||||
grouped[s['category']].append(s)
|
||||
|
||||
for cat in grouped:
|
||||
grouped[cat].sort(key=lambda x: x['confidence'], reverse=True)
|
||||
grouped[cat] = grouped[cat][:top_k_per_category]
|
||||
|
||||
return grouped
|
||||
Reference in New Issue
Block a user