bvandeusen
8af9f12544
PIL's strict load was raising OSError on images missing the trailing end-of-stream marker (e.g. '6 bytes not processed' on a JPEG without its FF D9 EOI), failing the entire ML task for an image the model could otherwise score fine. Set ImageFile.LOAD_TRUNCATED_IMAGES = True in both ML modules so a minutely-corrupt tail doesn't block tagging or embedding. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
117 lines
4.1 KiB
Python
117 lines
4.1 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, ImageFile
|
|
|
|
# Some images in the library are minutely truncated (e.g. 6 bytes short of
|
|
# the JPEG end-of-image marker). The model doesn't care about the last few
|
|
# pixels, but PIL's default strict load raises OSError. Tolerate them so a
|
|
# single corrupt tail doesn't block tagging the rest of the image.
|
|
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
|
|
|
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]
|