322 lines
11 KiB
Python
322 lines
11 KiB
Python
# -*- 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", "")
|
||
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)})
|
||
|
||
# 只返回评分，不做阈值判定
|
||
# WD Tagger v3 的 rating 是独立 sigmoid（不归一化），四个值可以同时很高。
|
||
|
||
ratings = result["ratings"]
|
||
top_tags = result["top_tags"]
|
||
explicit = ratings["explicit"]
|
||
questionable = ratings["questionable"]
|
||
|
||
# 打印 top_tags 到日志（调试用）
|
||
if top_tags:
|
||
tag_str = ", ".join([f"{t['tag']}:{t['score']:.2f}" for t in top_tags[:5]])
|
||
print(f"=== AnimeNSFW 标签: {tag_str}", flush=True)
|
||
|
||
# nsfw_score 返回 max(explicit, questionable) 的最高分
|
||
nsfw_score = max(explicit, questionable)
|
||
|
||
self._send({
|
||
"ok": True,
|
||
"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()
|