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

- 新增 anime_nsfw 模块:WD Tagger v3 模型检测动漫/二次元图片

- 修复误判:当 general 评分最高时判定为安全,避免非动漫图误判

- 数据库存储四个 rating 原始分数(JSON),无需重扫即可查看

- 媒体库新增动漫色情检测列、文件路径列(支持一键复制)

- 修复预览图加载占位符显示两个点的问题

- 移动文件按钮不再限制扫描完成后才显示
This commit is contained in:
2025-08-24 21:14:04 +08:00
parent bd84bc98e8
commit d6798ef1ee
12 changed files with 934 additions and 31 deletions
+335
View File
@@ -0,0 +1,335 @@
// Package anime_nsfw 负责"动漫色情图片识别":本机 Python + WD Tagger v3 (eva02-large) 起一个本地 HTTP 服务,
// Go 侧把图片路径发给它,拿到"是否色情"和最高敏感分。
//
// 大白话理解:WD Tagger v3 是 Danbooru 风格的图片标签模型,能输出 rating 分类:
// - general:安全
// - sensitive:轻微敏感
// - questionable:可疑
// - explicit:明确色情
//
// 我们只关心 questionable 和 explicit 的分数,任一 >= 阈值就判定为色情。
//
// 机器上没有 python 或没装 onnxruntime/opencv 时,Start 会报错,调用方跳过识别、
// 走默认整理流程——这个功能是纯附加的,不影响其它步骤。
package anime_nsfw
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net"
"net/http"
"os"
"os/exec"
"path/filepath"
"strings"
"sync"
"time"
_ "embed"
)
//go:embed server.py
var serverScript []byte
// detectorStartTimeout 等待 python 服务就绪(模型加载)的超时时间。
// WD Tagger v3 模型约 1.26GB,加载需要时间,给足余量。
const detectorStartTimeout = 300 * time.Second
// requestTimeout 单张图片检测的超时时间。
const requestTimeout = 60 * time.Second
// Detector 封装一个本地 python 服务的客户端。
type Detector struct {
url string
threshold float64
client *http.Client
cmd *exec.Cmd
scriptDir string
python string
modelDir string
mu sync.Mutex
stopped bool
cancel context.CancelFunc
done chan struct{}
restartMu sync.Mutex
}
// Start 拉起本地 python 检测服务并等它就绪。
func Start(pythonPath string, threshold float64) (*Detector, error) {
if threshold <= 0 || threshold >= 1 {
threshold = 0.5
}
py := resolvePython(pythonPath)
if py == "" {
return nil, fmt.Errorf("未找到 python,动漫敏感识别不可用")
}
tmpDir, err := os.MkdirTemp("", "kmp-anime-nsfw-*")
if err != nil {
return nil, fmt.Errorf("创建临时目录失败: %v", err)
}
scriptPath := filepath.Join(tmpDir, "server.py")
if err := os.WriteFile(scriptPath, serverScript, 0644); err != nil {
_ = os.RemoveAll(tmpDir)
return nil, fmt.Errorf("写入检测脚本失败: %v", err)
}
modelDir := ""
if wd, err := os.Getwd(); err == nil {
modelDir = filepath.Join(wd, "lib")
}
d := &Detector{
threshold: threshold,
client: &http.Client{Timeout: requestTimeout},
python: py,
modelDir: modelDir,
scriptDir: tmpDir,
}
if err := d.launch(); err != nil {
_ = os.RemoveAll(tmpDir)
return nil, err
}
return d, nil
}
func (d *Detector) launch() error {
d.teardown()
port, err := freePort()
if err != nil {
return fmt.Errorf("找空闲端口失败: %v", err)
}
ctx, cancel := context.WithCancel(context.Background())
cmd := exec.CommandContext(ctx, d.python, filepath.Join(d.scriptDir, "server.py"), fmt.Sprintf("%d", port), d.modelDir)
cmd.Stdout = os.Stdout
cmd.Stderr = os.Stdout
if err := cmd.Start(); err != nil {
cancel()
return fmt.Errorf("启动动漫检测服务失败: %v", err)
}
done := make(chan struct{})
d.url = fmt.Sprintf("http://127.0.0.1:%d", port)
d.cmd = cmd
d.cancel = cancel
d.done = done
go func() {
_ = cmd.Wait()
close(done)
}()
if err := d.waitReady(); err != nil {
d.teardown()
return err
}
return nil
}
func (d *Detector) teardown() {
if d.cancel != nil {
d.cancel()
d.cancel = nil
}
if d.cmd != nil && d.cmd.Process != nil {
_ = d.cmd.Process.Kill()
}
if d.done != nil {
select {
case <-d.done:
case <-time.After(3 * time.Second):
}
d.done = nil
}
d.cmd = nil
}
func (d *Detector) waitReady() error {
deadline := time.Now().Add(detectorStartTimeout)
for time.Now().Before(deadline) {
select {
case <-d.done:
return fmt.Errorf("动漫检测服务提前退出,请确认已安装 onnxruntime/opencv (pip install onnxruntime opencv-python numpy)")
default:
}
req, err := http.NewRequest(http.MethodGet, d.url+"/health", nil)
if err != nil {
return err
}
resp, err := d.client.Do(req)
if err == nil {
_, _ = io.Copy(io.Discard, resp.Body)
resp.Body.Close()
if resp.StatusCode == http.StatusOK {
return nil
}
}
time.Sleep(500 * time.Millisecond)
}
return fmt.Errorf("动漫检测服务在 %v 内未就绪", detectorStartTimeout)
}
// AnimeRatings WD Tagger 的四个 rating 评分。
type AnimeRatings struct {
General float64 `json:"general"`
Sensitive float64 `json:"sensitive"`
Questionable float64 `json:"questionable"`
Explicit float64 `json:"explicit"`
}
// Detect 检测一张动漫图片是否色情。返回:
// - score:max(explicit, questionable) 的 sigmoid 分数(若 G 最高则为 0)
// - nsfw:是否判定为色情(分数 >= 阈值)
// - ratings:四个 rating 的原始分数
// - err:检测失败
func (d *Detector) Detect(path string) (score float64, nsfw bool, ratings AnimeRatings, err error) {
score, nsfw, ratings, err = d.detectOnce(path)
if err == nil {
return score, nsfw, ratings, nil
}
var ce *connError
if !errors.As(err, &ce) {
return score, nsfw, ratings, err
}
d.restartMu.Lock()
defer d.restartMu.Unlock()
if d.isStopped() {
return 0, false, ratings, err
}
if d.healthy() {
return d.detectOnce(path)
}
fmt.Printf("=== 动漫检测服务失联(%v), 正在重启…\n", ce.err)
if rerr := d.launch(); rerr != nil {
return 0, false, ratings, fmt.Errorf("动漫检测服务重启失败: %v", rerr)
}
fmt.Println("=== 动漫检测服务已重启,继续识别")
return d.detectOnce(path)
}
func (d *Detector) healthy() bool {
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodGet, d.url+"/health", nil)
if err != nil {
return false
}
resp, err := d.client.Do(req)
if err != nil {
return false
}
defer resp.Body.Close()
return resp.StatusCode == http.StatusOK
}
func (d *Detector) isStopped() bool {
d.mu.Lock()
defer d.mu.Unlock()
return d.stopped
}
func (d *Detector) detectOnce(path string) (score float64, nsfw bool, ratings AnimeRatings, err error) {
body, _ := json.Marshal(map[string]interface{}{
"path": path,
"threshold": d.threshold,
})
resp, err := d.client.Post(d.url+"/detect", "application/json", bytes.NewReader(body))
if err != nil {
return 0, false, ratings, &connError{err}
}
defer resp.Body.Close()
data, err := io.ReadAll(resp.Body)
if err != nil {
return 0, false, ratings, err
}
var out struct {
Ok bool `json:"ok"`
Error string `json:"error"`
Nsfw bool `json:"nsfw"`
Score float64 `json:"score"`
Ratings struct {
General float64 `json:"general"`
Sensitive float64 `json:"sensitive"`
Questionable float64 `json:"questionable"`
Explicit float64 `json:"explicit"`
} `json:"ratings"`
TopTags []map[string]interface{} `json:"top_tags"`
}
if err := json.Unmarshal(data, &out); err != nil {
return 0, false, ratings, fmt.Errorf("解析检测结果失败: %v", err)
}
if !out.Ok {
return 0, false, ratings, fmt.Errorf("检测服务返回错误: %s", out.Error)
}
ratings = AnimeRatings{
General: out.Ratings.General,
Sensitive: out.Ratings.Sensitive,
Questionable: out.Ratings.Questionable,
Explicit: out.Ratings.Explicit,
}
// 打印 rating 详情(调试用)
fmt.Printf("=== AnimeNSFW 检测详情 %s: [G=%.3f S=%.3f Q=%.3f E=%.3f]\n",
path, ratings.General, ratings.Sensitive, ratings.Questionable, ratings.Explicit)
return out.Score, out.Nsfw, ratings, nil
}
type connError struct{ err error }
func (e *connError) Error() string { return e.err.Error() }
func (e *connError) Unwrap() error { return e.err }
func (d *Detector) Stop() {
d.mu.Lock()
defer d.mu.Unlock()
if d.stopped {
return
}
d.stopped = true
d.teardown()
if d.scriptDir != "" {
_ = os.RemoveAll(d.scriptDir)
}
}
func (d *Detector) Threshold() float64 {
return d.threshold
}
func resolvePython(configured string) string {
candidates := []string{}
if strings.TrimSpace(configured) != "" {
candidates = append(candidates, strings.TrimSpace(configured))
}
if v := os.Getenv("KM_NSFW_PYTHON"); v != "" {
candidates = append(candidates, v)
}
candidates = append(candidates, "python", "python3", "py")
for _, c := range candidates {
if p, err := exec.LookPath(c); err == nil {
return p
}
if strings.ContainsAny(c, `\/`) {
if _, err := os.Stat(c); err == nil {
return c
}
}
}
return ""
}
func freePort() (int, error) {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
return 0, err
}
defer ln.Close()
return ln.Addr().(*net.TCPAddr).Port, nil
}
+319
View File
@@ -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()