feat: 添加动漫色情检测(WD Tagger v3)，媒体库展示优化

- 新增 anime_nsfw 模块：WD Tagger v3 模型检测动漫/二次元图片 - 修复误判：当 general 评分最高时判定为安全，避免非动漫图误判 - 数据库存储四个 rating 原始分数(JSON)，无需重扫即可查看 - 媒体库新增动漫色情检测列、文件路径列(支持一键复制) - 修复预览图加载占位符显示两个点的问题 - 移动文件按钮不再限制扫描完成后才显示
2025-08-24 21:14:04 +08:00
parent bd84bc98e8
commit d6798ef1ee
12 changed files with 934 additions and 31 deletions
@@ -0,0 +1,319 @@
+# -*- coding: utf-8 -*-
+"""WD Tagger v3 (eva02-large) 动漫色情检测服务：由 Go 侧启动，常驻提供 /health 和 /detect。
+
+只依赖 Python 标准库 + onnxruntime + opencv-python + numpy。
+onnxruntime 的 InferenceSession 不是线程安全的，同一个实例并发推理会崩，
+所以这里用一个固定大小的推理线程池：每个 worker 线程各持一份独立模型实例，
+N 个请求可以真·并行推理。
+
+模型来源：SmilingWolf/wd-eva02-large-tagger-v3 (HuggingFace)
+  - model.onnx：ONNX 模型（约 1.26GB）
+  - selected_tags.csv：标签映射表，前 4 行是 rating 分类：
+      index 0: general（安全）
+      index 1: sensitive（轻微敏感）
+      index 2: questionable（可疑）
+      index 3: explicit（明确色情）
+
+模型文件放在 <程序运行目录>/lib/anime_nsfw/ 下，首次启动自动从 HuggingFace 下载。
+
+输入预处理：
+  - 图片 pad 成正方形（白边填充）
+  - bicubic 缩放到 448x448
+  - 归一化：(pixel/255 - 0.5) / 0.5
+
+NSFW 判定逻辑：
+  - explicit 或 questionable 的 sigmoid 分数 >= 阈值 → 判定为色情
+  - score 返回 max(explicit, questionable) 的 sigmoid 分数
+"""
+import concurrent.futures
+import json
+import os
+import shutil
+import sys
+import threading
+import urllib.request
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+import cv2
+import numpy as np
+import onnxruntime as ort
+
+# 模型文件与 HuggingFace 下载地址
+HF_REPO = "SmilingWolf/wd-eva02-large-tagger-v3"
+HF_BASE_URL = "https://huggingface.co/%s/resolve/main" % HF_REPO
+MODEL_FILENAME = "model.onnx"
+TAGS_FILENAME = "selected_tags.csv"
+
+# 模型输入参数（来自 config.json）
+INPUT_SIZE = 448
+MEAN = np.array([0.5, 0.5, 0.5], dtype=np.float32)
+STD = np.array([0.5, 0.5, 0.5], dtype=np.float32)
+
+# Rating 标签在 selected_tags.csv 中的位置（前 4 行）
+RATING_GENERAL = 0
+RATING_SENSITIVE = 1
+RATING_QUESTIONABLE = 2
+RATING_EXPLICIT = 3
+
+# 模型目录由 Go 侧作为第二个参数传入（程序运行目录下的 lib）
+MODEL_DIR = sys.argv[2] if len(sys.argv) > 2 else ""
+
+# 推理并发数（每个 worker 持独立模型实例）
+INFER_WORKERS = int(os.environ.get("KM_ANIME_NSFW_WORKERS", "2") or "2")
+if INFER_WORKERS < 1:
+    INFER_WORKERS = 1
+
+_executor = None
+_thread_local = threading.local()
+_MODEL_PATH = None
+_TAGS = None  # numpy array of tag names (loaded once)
+
+
+def _looks_like_model(path):
+    """粗略校验文件是不是真 ONNX 模型。"""
+    try:
+        with open(path, "rb") as f:
+            head = f.read(4)
+        if not head or head.startswith(b"<"):
+            return False
+        return os.path.getsize(path) > 100_000_000  # ONNX 模型约 1.26GB
+    except OSError:
+        return False
+
+
+def _download_file(url, local):
+    """带超时地把文件下载到 local。"""
+    req = urllib.request.Request(url)
+    with urllib.request.urlopen(req, timeout=120) as resp, open(local, "wb") as f:
+        shutil.copyfileobj(resp, f)
+
+
+def _ensure_model(model_dir):
+    """确保模型和标签文件都存在，返回 (model_path, tags_csv_path)。"""
+    if not model_dir:
+        raise RuntimeError("未指定模型目录")
+
+    anime_dir = os.path.join(model_dir, "anime_nsfw")
+    os.makedirs(anime_dir, exist_ok=True)
+
+    model_path = os.path.join(anime_dir, MODEL_FILENAME)
+    tags_path = os.path.join(anime_dir, TAGS_FILENAME)
+
+    # 下载模型（如果本地没有或无效）
+    if not os.path.isfile(model_path) or not _looks_like_model(model_path):
+        url = "%s/%s" % (HF_BASE_URL, MODEL_FILENAME)
+        print("=== [AnimeNSFW] 正在下载模型: %s" % url, flush=True)
+        _download_file(url, model_path)
+        if not _looks_like_model(model_path):
+            raise RuntimeError("下载的模型文件无效")
+        print("=== [AnimeNSFW] 模型下载完成: %s" % model_path, flush=True)
+    else:
+        print("=== [AnimeNSFW] 使用本地模型: %s" % model_path, flush=True)
+
+    # 下载标签 CSV（如果本地没有）
+    if not os.path.isfile(tags_path):
+        url = "%s/%s" % (HF_BASE_URL, TAGS_FILENAME)
+        print("=== [AnimeNSFW] 正在下载标签: %s" % url, flush=True)
+        _download_file(url, tags_path)
+        print("=== [AnimeNSFW] 标签下载完成: %s" % tags_path, flush=True)
+
+    return model_path, tags_path
+
+
+def _load_tags(tags_path):
+    """加载 selected_tags.csv，返回 tag name 的 numpy array。"""
+    import csv
+    tags = []
+    with open(tags_path, "r", encoding="utf-8") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            tags.append(row["name"])
+    return np.array(tags)
+
+
+def _new_session(model_path):
+    """创建一份独立 ONNX 推理会话（每个 worker 各持一份）。"""
+    sess_opts = ort.SessionOptions()
+    sess_opts.inter_op_num_threads = 1
+    return ort.InferenceSession(model_path, sess_options=sess_opts, providers=["CPUExecutionProvider"])
+
+
+def _preprocess(data):
+    """图片字节 → 模型输入 tensor。
+
+    步骤：
+      1. 解码图片（从字节）
+      2. RGBA → BGR
+      3. Pad 成正方形（白边 255 填充）
+      4. Bicubic 缩放到 INPUT_SIZE x INPUT_SIZE
+      5. 归一化：(pixel/255 - mean) / std
+      6. HWC → NCHW
+    """
+    mat = cv2.imdecode(np.frombuffer(data, np.uint8), -1)
+    if mat is None:
+        raise OSError("图片解码失败")
+
+    # RGBA → BGR
+    if mat.ndim == 3 and mat.shape[2] == 4:
+        mat = cv2.cvtColor(mat, cv2.COLOR_RGBA2BGR)
+    elif mat.ndim == 2:
+        mat = cv2.cvtColor(mat, cv2.COLOR_GRAY2BGR)
+
+    # Pad 成正方形
+    h, w = mat.shape[:2]
+    max_size = max(w, h)
+    x_pad = max_size - w
+    y_pad = max_size - h
+    mat_pad = cv2.copyMakeBorder(mat, 0, y_pad, 0, x_pad, cv2.BORDER_CONSTANT, value=(255, 255, 255))
+
+    # Bicubic 缩放到模型输入尺寸
+    resized = cv2.resize(mat_pad, (INPUT_SIZE, INPUT_SIZE), interpolation=cv2.INTER_CUBIC)
+
+    # 归一化：BGR → RGB, uint8 → float32, (pixel/255 - mean) / std
+    img = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
+    img = (img - MEAN) / STD
+
+    # 模型输入格式 NHWC → NCHW 不需要，直接加 batch 维度
+    img = np.expand_dims(img, axis=0)
+    return img
+
+
+def _sigmoid(x):
+    """数值稳定的 sigmoid。"""
+    return np.where(x >= 0, 1.0 / (1.0 + np.exp(-x)), np.exp(x) / (1.0 + np.exp(x)))
+
+
+def _infer(path):
+    """整个推理任务（读文件 + 预处理 + 推理 + 后处理）。"""
+    # 读文件
+    lp = path
+    if os.name == "nt":
+        if not lp.startswith("\\\\?\\"):
+            if lp.startswith("\\\\"):
+                lp = "\\\\?\\UNC\\" + lp.lstrip("\\")
+            elif os.path.isabs(lp):
+                lp = "\\\\?\\" + os.path.normpath(lp)
+    if not os.path.exists(lp):
+        raise OSError("文件不存在: %s" % path)
+    with open(lp, "rb") as f:
+        data = f.read()
+    if not data:
+        raise OSError("空文件: %s" % path)
+
+    # 预处理
+    input_tensor = _preprocess(data)
+
+    # 推理（每个 worker 线程持独立 session）
+    sess = getattr(_thread_local, "session", None)
+    if sess is None:
+        sess = _new_session(_MODEL_PATH)
+        _thread_local.session = sess
+
+    input_name = sess.get_inputs()[0].name
+    outputs = sess.run(None, {input_name: input_tensor})
+    logits = outputs[0][0]  # shape: (num_classes,)
+
+    # Sigmoid 转概率
+    probs = _sigmoid(logits)
+
+    # 提取 rating 分数
+    rating_general = float(probs[RATING_GENERAL])
+    rating_sensitive = float(probs[RATING_SENSITIVE])
+    rating_questionable = float(probs[RATING_QUESTIONABLE])
+    rating_explicit = float(probs[RATING_EXPLICIT])
+
+    # 找 top 10 一般标签（跳过前 4 个 rating）
+    tag_probs = probs[4:]
+    top_indices = np.argsort(tag_probs)[::-1][:10]
+    top_tags = []
+    for idx in top_indices:
+        tag_name = str(_TAGS[idx + 4])
+        tag_score = float(tag_probs[idx])
+        if tag_score > 0.05:
+            top_tags.append({"tag": tag_name, "score": round(tag_score, 4)})
+
+    return {
+        "ratings": {
+            "general": round(rating_general, 4),
+            "sensitive": round(rating_sensitive, 4),
+            "questionable": round(rating_questionable, 4),
+            "explicit": round(rating_explicit, 4),
+        },
+        "top_tags": top_tags,
+    }
+
+
+class Handler(BaseHTTPRequestHandler):
+    def _send(self, obj, status=200):
+        data = json.dumps(obj, ensure_ascii=False).encode("utf-8")
+        self.send_response(status)
+        self.send_header("Content-Type", "application/json; charset=utf-8")
+        self.send_header("Content-Length", str(len(data)))
+        self.end_headers()
+        self.wfile.write(data)
+
+    def do_GET(self):
+        if self.path == "/health":
+            self._send({"ok": True})
+        else:
+            self._send({"ok": False, "error": "not found"}, 404)
+
+    def do_POST(self):
+        if self.path != "/detect":
+            return self._send({"ok": False, "error": "not found"}, 404)
+        length = int(self.headers.get("Content-Length", 0))
+        try:
+            body = json.loads(self.rfile.read(length) or b"{}")
+            path = body.get("path", "")
+            threshold = float(body.get("threshold", 0.5))
+            if not path:
+                return self._send({"ok": False, "error": "缺少图片路径"})
+            result = _executor.submit(_infer, path).result()
+        except Exception as exc:
+            return self._send({"ok": False, "error": str(exc)})
+
+        # NSFW 判定：explicit 或 questionable 分数 >= 阈值，且必须高于 general。
+        # WD Tagger 的 rating 是独立 sigmoid（不归一化），非动漫图四个分值往往
+        # 都很接近；若 general 最高说明模型认为是安全图，只看绝对值会大量误判。
+        ratings = result["ratings"]
+        nsfw_raw = max(ratings["explicit"], ratings["questionable"])
+        if nsfw_raw > ratings["general"]:
+            nsfw_score = nsfw_raw
+            nsfw = nsfw_score >= threshold
+        else:
+            nsfw_score = 0.0
+            nsfw = False
+
+        self._send({
+            "ok": True,
+            "nsfw": nsfw,
+            "score": nsfw_score,
+            "ratings": ratings,
+            "top_tags": result["top_tags"],
+        })
+
+    def log_message(self, *args):
+        pass
+
+
+class Server(ThreadingHTTPServer):
+    request_queue_size = 128
+
+
+if __name__ == "__main__":
+    port = int(sys.argv[1]) if len(sys.argv) > 1 else 0
+
+    # 加载模型和标签
+    _MODEL_PATH, tags_path = _ensure_model(MODEL_DIR)
+    _TAGS = _load_tags(tags_path)
+    print("=== [AnimeNSFW] 标签加载完成, 共 %d 个" % len(_TAGS), flush=True)
+
+    # 预创建第一个 session 验证模型可用
+    sess = _new_session(_MODEL_PATH)
+    print("=== [AnimeNSFW] 模型加载成功, 输入: %s" % str(sess.get_inputs()[0].shape), flush=True)
+    del sess
+
+    # 启动推理线程池
+    _executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS)
+    print("=== [AnimeNSFW] 服务启动, 端口=%d, workers=%d" % (port, INFER_WORKERS), flush=True)
+    Server(("127.0.0.1", port), Handler).serve_forever()