357 lines
12 KiB
Python
357 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""wd-tagger 本地检测服务：由 Go 侧启动，常驻提供 /health 和 /detect。
|
||
|
||
基于 SmilingWolf/wd-vit-large-tagger-v3 模型，专门用于二次元（anime/hentai）图片分级。
|
||
输出 rating:safe / rating:questionable / rating:explicit 三个分级分数，
|
||
以及若干高置信度标签（用于辅助判断内容类型）。
|
||
|
||
只依赖 Python 标准库 + onnxruntime + Pillow。
|
||
模型与标签文件放在 <程序运行目录>/lib/wd-tagger/ 下，首次启动自动从 HuggingFace 下载。
|
||
|
||
推理并发策略与 nsfw/server.py 相同：每个 worker 线程各持一份独立 session 实例，
|
||
避免 onnxruntime 同一实例并发推理崩溃的问题。
|
||
"""
|
||
import concurrent.futures
|
||
import csv
|
||
import io
|
||
import json
|
||
import os
|
||
import shutil
|
||
import sys
|
||
import tempfile
|
||
import threading
|
||
import urllib.request
|
||
import zipfile
|
||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||
|
||
try:
|
||
import numpy as np
|
||
except ImportError:
|
||
print("=== 错误: 缺少 numpy，请执行 pip install numpy", flush=True)
|
||
sys.exit(1)
|
||
|
||
try:
|
||
import onnxruntime as ort
|
||
except ImportError:
|
||
print("=== 错误: 缺少 onnxruntime，请执行 pip install onnxruntime", flush=True)
|
||
sys.exit(1)
|
||
|
||
try:
|
||
from PIL import Image
|
||
except ImportError:
|
||
print("=== 错误: 缺少 Pillow，请执行 pip install Pillow", flush=True)
|
||
sys.exit(1)
|
||
|
||
# 模型相关配置
|
||
MODEL_DIR_NAME = "wd-tagger"
|
||
MODEL_FILENAME = "model.onnx"
|
||
TAGS_FILENAME = "selected_tags.csv"
|
||
|
||
# HuggingFace 下载地址（wd-vit-tagger-v3，轻量版，约 340MB）
|
||
HF_MODEL_URL = "https://huggingface.co/SmilingWolf/wd-vit-tagger-v3/resolve/main/model.onnx"
|
||
HF_TAGS_URL = "https://huggingface.co/SmilingWolf/wd-vit-tagger-v3/resolve/main/selected_tags.csv"
|
||
|
||
# 模型输入尺寸
|
||
INPUT_SIZE = 448
|
||
|
||
# ImageNet 归一化参数
|
||
IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
|
||
IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
|
||
|
||
# 我们关注的分级标签
|
||
RATING_TAGS = {
|
||
"rating:safe": "safe",
|
||
"rating:questionable": "questionable",
|
||
"rating:explicit": "explicit",
|
||
}
|
||
|
||
# 模型目录由 Go 侧作为第二个参数传入
|
||
MODEL_DIR = sys.argv[2] if len(sys.argv) > 2 else ""
|
||
|
||
# 推理并发数（ViT-Large 模型较大，默认 2 避免 CPU 争抢）
|
||
INFER_WORKERS = int(os.environ.get("KM_HENTAI_WORKERS", "2") or "2")
|
||
if INFER_WORKERS < 1:
|
||
INFER_WORKERS = 1
|
||
|
||
_executor = None
|
||
_thread_local = threading.local()
|
||
|
||
# 标签表：启动时从 CSV 加载
|
||
_tag_names = []
|
||
_rating_indices = {} # tag名 -> 在输出数组中的索引
|
||
|
||
|
||
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) > 1_000_000
|
||
except OSError:
|
||
return False
|
||
|
||
|
||
def _download_file(url, local):
|
||
"""带超时地下载文件。"""
|
||
req = urllib.request.Request(url)
|
||
with urllib.request.urlopen(req, timeout=60) as resp, open(local, "wb") as f:
|
||
shutil.copyfileobj(resp, f)
|
||
|
||
|
||
def _ensure_model(model_dir):
|
||
"""确保模型和标签文件就绪，返回 (model_path, tags_path) 或 (None, None) 表示用包自带。"""
|
||
if not model_dir:
|
||
return None, None
|
||
|
||
wd_dir = os.path.join(model_dir, MODEL_DIR_NAME)
|
||
model_path = os.path.join(wd_dir, MODEL_FILENAME)
|
||
tags_path = os.path.join(wd_dir, TAGS_FILENAME)
|
||
|
||
# 检查模型文件
|
||
if os.path.isfile(model_path) and _looks_like_model(model_path):
|
||
print("=== 使用本地 wd-tagger 模型: %s" % model_path, flush=True)
|
||
else:
|
||
# 下载模型
|
||
try:
|
||
os.makedirs(wd_dir, exist_ok=True)
|
||
print("=== 本地未找到 wd-tagger 模型, 正在下载: %s" % model_path, flush=True)
|
||
_download_file(HF_MODEL_URL, model_path)
|
||
if _looks_like_model(model_path):
|
||
print("=== wd-tagger 模型下载完成: %s" % model_path, flush=True)
|
||
else:
|
||
print("=== 下载的模型文件无效, 删除", flush=True)
|
||
try:
|
||
os.remove(model_path)
|
||
except OSError:
|
||
pass
|
||
return None, None
|
||
except Exception as exc:
|
||
print("=== wd-tagger 模型下载失败: %s" % exc, flush=True)
|
||
return None, None
|
||
|
||
# 检查标签文件
|
||
if os.path.isfile(tags_path):
|
||
print("=== 使用本地标签文件: %s" % tags_path, flush=True)
|
||
else:
|
||
try:
|
||
print("=== 下载标签文件: %s" % tags_path, flush=True)
|
||
_download_file(HF_TAGS_URL, tags_path)
|
||
print("=== 标签文件下载完成", flush=True)
|
||
except Exception as exc:
|
||
# 标签文件下载失败不影响模型运行，只是无法解析输出标签
|
||
print("=== 警告: 标签文件下载失败，将无法解析输出标签: %s" % exc, flush=True)
|
||
tags_path = "" # 标记标签文件不可用，但仍返回模型路径
|
||
|
||
return model_path, tags_path
|
||
|
||
|
||
def _load_tags(tags_path):
|
||
"""加载标签 CSV，建立 tag名 -> 索引 的映射。"""
|
||
global _tag_names, _rating_indices
|
||
_tag_names = []
|
||
_rating_indices = {}
|
||
|
||
if not tags_path or not os.path.isfile(tags_path):
|
||
print("=== 警告: 标签文件不存在，无法解析输出", flush=True)
|
||
return
|
||
|
||
with open(tags_path, "r", encoding="utf-8") as f:
|
||
reader = csv.DictReader(f)
|
||
for row in reader:
|
||
name = row.get("name", "")
|
||
_tag_names.append(name)
|
||
if name in RATING_TAGS:
|
||
_rating_indices[name] = len(_tag_names) - 1
|
||
|
||
print("=== 已加载 %d 个标签, 其中分级标签: %s" % (len(_tag_names), list(_rating_indices.keys())), flush=True)
|
||
|
||
|
||
def _new_session(model_path):
|
||
"""创建一份独立的 ONNX 推理 session。"""
|
||
if model_path:
|
||
return ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
|
||
# 没有本地模型，无法推理
|
||
raise RuntimeError("wd-tagger 模型不可用")
|
||
|
||
|
||
def get_session():
|
||
"""确保模型已就绪，返回一个推理 session（启动时调用）。"""
|
||
global _session
|
||
if "_session" not in dir() or _session is None:
|
||
model_path, tags_path = _ensure_model(MODEL_DIR)
|
||
if not model_path:
|
||
raise RuntimeError("wd-tagger 模型不可用")
|
||
if not tags_path:
|
||
raise RuntimeError("wd-tagger 标签文件不可用，无法解析输出。请手动下载 %s 到 %s" % (TAGS_FILENAME, os.path.join(MODEL_DIR, MODEL_DIR_NAME)))
|
||
_load_tags(tags_path)
|
||
_session = _new_session(model_path)
|
||
return _session
|
||
|
||
|
||
_session = None
|
||
|
||
|
||
def _preprocess(img_bytes):
|
||
"""把图片字节预处理成模型输入的 numpy 数组。
|
||
|
||
步骤：
|
||
1. 解码为 PIL Image
|
||
2. 转为 RGB（如果有 alpha 通道）
|
||
3. 直接 resize 到 INPUT_SIZE x INPUT_SIZE（官方推荐方式）
|
||
4. 归一化（ImageNet 参数）
|
||
5. 转为 NHWC 格式的 float32 数组（模型期望 [batch, H, W, C]）
|
||
"""
|
||
img = Image.open(io.BytesIO(img_bytes))
|
||
if img.mode != "RGB":
|
||
img = img.convert("RGB")
|
||
|
||
# 直接 resize 到目标尺寸（wd-tagger 官方推荐）
|
||
img = img.resize((INPUT_SIZE, INPUT_SIZE), Image.LANCZOS)
|
||
|
||
# 转 numpy 数组并归一化
|
||
arr = np.array(img, dtype=np.float32) / 255.0
|
||
arr = (arr - IMAGENET_MEAN) / IMAGENET_STD
|
||
|
||
# HWC -> NHWC (模型期望 [batch, height, width, channels])
|
||
arr = np.expand_dims(arr, axis=0)
|
||
return arr
|
||
|
||
|
||
def _infer_bytes(session, img_bytes):
|
||
"""用给定的 session 推理一张图片（字节），返回 (ratings_dict, top_tags_list)。"""
|
||
inp = _preprocess(img_bytes)
|
||
input_name = session.get_inputs()[0].name
|
||
outputs = session.run(None, {input_name: inp})
|
||
|
||
# 输出是 [1, num_tags] 的分数数组
|
||
scores = outputs[0][0] if len(outputs[0].shape) == 2 else outputs[0]
|
||
|
||
# 提取分级分数
|
||
ratings = {}
|
||
for tag_name, idx in _rating_indices.items():
|
||
if idx < len(scores):
|
||
rating_key = RATING_TAGS[tag_name]
|
||
ratings[rating_key] = float(scores[idx])
|
||
|
||
# 提取 top 标签（排除分级标签，取分数最高的 10 个）
|
||
top_tags = []
|
||
non_rating_indices = [i for i in range(len(scores)) if i not in _rating_indices.values()]
|
||
if non_rating_indices:
|
||
top_indices = sorted(non_rating_indices, key=lambda i: scores[i], reverse=True)[:10]
|
||
for idx in top_indices:
|
||
if idx < len(_tag_names) and scores[idx] > 0.5:
|
||
top_tags.append({"tag": _tag_names[idx], "score": float(scores[idx])})
|
||
|
||
return ratings, top_tags
|
||
|
||
|
||
def _read_bytes(path):
|
||
"""用 Python 读文件字节。"""
|
||
if os.name == "nt":
|
||
if path.startswith("\\\\?\\"):
|
||
lp = path
|
||
elif path.startswith("\\\\"):
|
||
lp = "\\\\?\\UNC\\" + path.lstrip("\\")
|
||
elif os.path.isabs(path):
|
||
lp = "\\\\?\\" + os.path.normpath(path)
|
||
else:
|
||
lp = path
|
||
else:
|
||
lp = path
|
||
|
||
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)
|
||
return data
|
||
|
||
|
||
def _infer(path):
|
||
"""整个推理任务：读文件 + 预处理 + 推理。"""
|
||
data = _read_bytes(path)
|
||
session = getattr(_thread_local, "session", None)
|
||
if session is None:
|
||
# 每个 worker 线程各持一份独立 session
|
||
model_path = os.path.join(MODEL_DIR, MODEL_DIR_NAME, MODEL_FILENAME) if MODEL_DIR else ""
|
||
session = _new_session(model_path)
|
||
_thread_local.session = session
|
||
return _infer_bytes(session, data)
|
||
|
||
|
||
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": "缺少图片路径"})
|
||
# 丢给推理线程池
|
||
ratings, top_tags = _executor.submit(_infer, path).result()
|
||
except Exception as exc:
|
||
return self._send({"ok": False, "error": str(exc)})
|
||
|
||
# 判断是否为 hentai：explicit 分数 >= 阈值 或 questionable 分数 >= 阈值
|
||
explicit_score = ratings.get("explicit", 0.0)
|
||
questionable_score = ratings.get("questionable", 0.0)
|
||
safe_score = ratings.get("safe", 0.0)
|
||
|
||
# hentai 判定：explicit 或 questionable 分数超过阈值
|
||
is_hentai = (explicit_score >= threshold) or (questionable_score >= threshold)
|
||
|
||
# 取最高的非 safe 分数作为"敏感分"
|
||
max_score = max(explicit_score, questionable_score)
|
||
|
||
# 确定分级标签
|
||
rating = "safe"
|
||
if explicit_score >= questionable_score and explicit_score >= threshold:
|
||
rating = "explicit"
|
||
elif questionable_score >= threshold:
|
||
rating = "questionable"
|
||
|
||
self._send({
|
||
"ok": True,
|
||
"hentai": is_hentai,
|
||
"score": max_score,
|
||
"rating": rating,
|
||
"ratings": ratings,
|
||
"topTags": 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
|
||
# 启动前先把模型准备好
|
||
get_session()
|
||
# 固定大小推理线程池
|
||
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS)
|
||
Server(("127.0.0.1", port), Handler).serve_forever()
|