Upload caption/wd14_tagger.py with huggingface_hub

caption/wd14_tagger.py

@@ -0,0 +1,151 @@
+#!/usr/bin/env python3
+"""WD14 tagger — danbooru-теги для аниме-картинок (SmilingWolf/wd-*-tagger-v3).
+
+Пишет теги в <имя_файла>.txt (danbooru-стиль: "1girl, animal ears, ...").
+Уже размеченные файлы пропускает (--force — перезаписать).
+
+Usage:
+    python wd14_tagger.py <папка_или_файл>... [--model vit] [--gen 0.5] [--char 0.85]
+
+Зависимости: torch, timm, pandas, pillow, huggingface_hub. Нужен GPU (или CPU).
+"""
+
+import argparse
+import time
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import timm
+import torch
+from huggingface_hub import hf_hub_download
+from huggingface_hub.utils import HfHubHTTPError
+from PIL import Image
+from timm.data import create_transform, resolve_data_config
+from torch import nn
+from torch.nn import functional as F
+from tqdm import tqdm
+
+MODEL_REPO_MAP = {
+    "vit": "SmilingWolf/wd-vit-tagger-v3",
+    "swinv2": "SmilingWolf/wd-swinv2-tagger-v3",
+    "convnext": "SmilingWolf/wd-convnext-tagger-v3",
+}
+
+IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
+
+
+def pil_ensure_rgb(image: Image.Image) -> Image.Image:
+    if image.mode not in ("RGB", "RGBA"):
+        image = image.convert("RGBA" if "transparency" in image.info else "RGB")
+    if image.mode == "RGBA":
+        canvas = Image.new("RGBA", image.size, (255, 255, 255))
+        canvas.alpha_composite(image)
+        image = canvas.convert("RGB")
+    return image
+
+
+def pil_pad_square(image: Image.Image) -> Image.Image:
+    w, h = image.size
+    px = max(image.size)
+    canvas = Image.new("RGB", (px, px), (255, 255, 255))
+    canvas.paste(image, ((px - w) // 2, (px - h) // 2))
+    return canvas
+
+
+@dataclass
+class LabelData:
+    names: list
+    rating: list
+    general: list
+    character: list
+
+
+def load_labels_hf(repo_id: str) -> LabelData:
+    try:
+        csv_path = hf_hub_download(repo_id=repo_id, filename="selected_tags.csv")
+    except HfHubHTTPError as e:
+        raise FileNotFoundError(f"selected_tags.csv failed to download from {repo_id}") from e
+    df = pd.read_csv(csv_path, usecols=["name", "category"])
+    return LabelData(
+        names=df["name"].tolist(),
+        rating=list(np.where(df["category"] == 9)[0]),
+        general=list(np.where(df["category"] == 0)[0]),
+        character=list(np.where(df["category"] == 4)[0]),
+    )
+
+
+def get_tags(probs, labels, gen_threshold, char_threshold):
+    """Probs -> строка тегов (danbooru-стиль, рейтинг в конце)."""
+    probs = list(zip(labels.names, probs.numpy()))
+    rating = dict(sorted((probs[i] for i in labels.rating), key=lambda x: x[1], reverse=True))
+    gen = {n: p for n, p in (probs[i] for i in labels.general) if p > gen_threshold}
+    char = {n: p for n, p in (probs[i] for i in labels.character) if p > char_threshold}
+    gen = dict(sorted(gen.items(), key=lambda x: x[1], reverse=True))
+    char = dict(sorted(char.items(), key=lambda x: x[1], reverse=True))
+    tags = list(gen) + list(char)
+    caption = ", ".join(t.replace("_", " ").replace("(", r"\(").replace(")", r"\)") for t in tags)
+    caption += ", rating_" + next(iter(rating), "")
+    return caption
+
+
+def main():
+    parser = argparse.ArgumentParser(description="WD14 tagger (danbooru-теги)")
+    parser.add_argument("paths", nargs="+", help="папки или файлы (папки обходятся рекурсивно)")
+    parser.add_argument("--model", default="vit", choices=sorted(MODEL_REPO_MAP))
+    parser.add_argument("--gen", dest="gen_threshold", type=float, default=0.5,
+                        help="порог уверенности для general-тегов")
+    parser.add_argument("--char", dest="char_threshold", type=float, default=0.85,
+                        help="порог уверенности для character-тегов")
+    parser.add_argument("--force", action="store_true", help="перезаписать существующие .txt")
+    args = parser.parse_args()
+
+    repo_id = MODEL_REPO_MAP[args.model]
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # собираем файлы
+    files = []
+    for arg in args.paths:
+        p = Path(arg)
+        if p.is_dir():
+            files.extend(f for f in p.rglob("*") if f.suffix.lower() in IMAGE_EXTS)
+        elif p.is_file():
+            files.append(p)
+    files = sorted(set(files))
+    if not files:
+        print("Нет файлов для обработки.")
+        return
+
+    print(f"Загрузка модели {args.model} ({repo_id}) на {device}...")
+    model = timm.create_model("hf-hub:" + repo_id).eval()
+    model.load_state_dict(timm.models.load_state_dict_from_hf(repo_id))
+    labels = load_labels_hf(repo_id)
+    transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
+    if device.type == "cuda":
+        model = model.to(device)
+
+    t0 = time.time()
+    done = skipped = 0
+    for image_path in tqdm(files, desc="Tagging", unit="img"):
+        txt_path = image_path.with_suffix(".txt")
+        if txt_path.exists() and not args.force:
+            skipped += 1
+            continue
+        try:
+            img = pil_pad_square(pil_ensure_rgb(Image.open(image_path)))
+            inputs = transform(img).unsqueeze(0)[:, [2, 1, 0]]  # RGB -> BGR
+            with torch.inference_mode():
+                outputs = F.sigmoid(model(inputs.to(device))).cpu()
+            caption = get_tags(outputs.squeeze(0), labels, args.gen_threshold, args.char_threshold)
+            txt_path.write_text(caption + "\n")
+            done += 1
+        except Exception as e:
+            tqdm.write(f"  SKIP {image_path.name}: {e}")
+
+    print(f"\nГотово: {done} размечено, {skipped} пропущено (уже есть .txt), "
+          f"за {time.strftime('%H:%M:%S', time.gmtime(time.time() - t0))}")
+
+
+if __name__ == "__main__":
+    main()