bvandeusen
0f35a0c484
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
111 lines
3.8 KiB
Python
111 lines
3.8 KiB
Python
"""WD14 EVA02-Large tagger (ONNX CPU inference)."""
|
|
from __future__ import annotations
|
|
import csv
|
|
import os
|
|
from typing import Iterable
|
|
|
|
import numpy as np
|
|
import onnxruntime as ort
|
|
from PIL import Image
|
|
|
|
MODEL_VERSION = os.environ.get('WD14_MODEL_VERSION', 'wd-eva02-large-tagger-v3')
|
|
_MODEL_DIR = os.environ.get('ML_MODEL_DIR', '/models')
|
|
_WD14_DIR = os.path.join(_MODEL_DIR, 'wd14')
|
|
_MODEL_PATH = os.path.join(_WD14_DIR, 'model.onnx')
|
|
_TAGS_PATH = os.path.join(_WD14_DIR, 'selected_tags.csv')
|
|
|
|
# WD14 selected_tags.csv uses Danbooru category ids:
|
|
# 0=general, 1=artist, 3=copyright, 4=character, 5=meta, 9=rating
|
|
_CATEGORY_MAP = {0: 'general', 1: 'artist', 3: 'copyright', 4: 'character', 5: 'meta', 9: 'rating'}
|
|
|
|
_session: ort.InferenceSession | None = None
|
|
_tag_meta: list[dict] | None = None
|
|
_input_name: str | None = None
|
|
_output_name: str | None = None
|
|
_input_size: int = 448
|
|
|
|
|
|
# NOT thread-safe. Must run in the ml-worker container with --concurrency=1.
|
|
def _load() -> None:
|
|
global _session, _tag_meta, _input_name, _output_name, _input_size
|
|
if _session is not None:
|
|
return
|
|
|
|
if not os.path.isfile(_MODEL_PATH):
|
|
raise RuntimeError(
|
|
f"WD14 model file missing at {_MODEL_PATH}. "
|
|
f"Populate the /models volume via the ml-worker downloader."
|
|
)
|
|
if not os.path.isfile(_TAGS_PATH):
|
|
raise RuntimeError(
|
|
f"WD14 selected_tags.csv missing at {_TAGS_PATH}. "
|
|
f"Populate the /models volume via the ml-worker downloader."
|
|
)
|
|
|
|
tag_meta: list[dict] = []
|
|
with open(_TAGS_PATH, newline='') as f:
|
|
reader = csv.DictReader(f)
|
|
for row in reader:
|
|
tag_meta.append({
|
|
'name': row['name'],
|
|
'category': _CATEGORY_MAP.get(int(row['category']), 'unknown'),
|
|
})
|
|
|
|
session = ort.InferenceSession(
|
|
_MODEL_PATH,
|
|
providers=['CPUExecutionProvider'],
|
|
)
|
|
_input_name = session.get_inputs()[0].name
|
|
_output_name = session.get_outputs()[0].name
|
|
# Input shape is usually [batch, H, W, 3] NHWC; pick the spatial dim
|
|
input_shape = session.get_inputs()[0].shape
|
|
for dim in input_shape:
|
|
if isinstance(dim, int) and dim > 1:
|
|
_input_size = dim
|
|
break
|
|
# Assign sentinels last so a partially-loaded state can't be observed.
|
|
_tag_meta = tag_meta
|
|
_session = session
|
|
|
|
|
|
def _preprocess(image_path: str) -> np.ndarray:
|
|
img = Image.open(image_path)
|
|
if img.mode != 'RGBA':
|
|
img = img.convert('RGBA')
|
|
# Composite onto white background so transparency doesn't bias the model
|
|
bg = Image.new('RGBA', img.size, (255, 255, 255, 255))
|
|
bg.paste(img, mask=img.split()[3] if img.mode == 'RGBA' else None)
|
|
img = bg.convert('RGB')
|
|
|
|
w, h = img.size
|
|
side = max(w, h)
|
|
square = Image.new('RGB', (side, side), (255, 255, 255))
|
|
square.paste(img, ((side - w) // 2, (side - h) // 2))
|
|
square = square.resize((_input_size, _input_size), Image.BICUBIC)
|
|
|
|
arr = np.array(square, dtype=np.float32)
|
|
# WD14 was trained on BGR
|
|
arr = arr[:, :, ::-1]
|
|
return arr[np.newaxis, :, :, :] # NHWC
|
|
|
|
|
|
def infer(image_path: str) -> list[dict]:
|
|
"""Run WD14 on one image. Returns a list of {name, category, confidence}."""
|
|
_load()
|
|
x = _preprocess(image_path)
|
|
out = _session.run([_output_name], {_input_name: x})[0][0]
|
|
results: list[dict] = []
|
|
for idx, score in enumerate(out):
|
|
meta = _tag_meta[idx]
|
|
results.append({
|
|
'name': meta['name'],
|
|
'category': meta['category'],
|
|
'confidence': float(score),
|
|
})
|
|
return results
|
|
|
|
|
|
def infer_filtered(image_path: str, min_any: float = 0.05) -> list[dict]:
|
|
"""Same as infer() but drops tags below a floor to keep DB rows reasonable."""
|
|
return [r for r in infer(image_path) if r['confidence'] >= min_any]
|