feat: 添加动漫色情检测(WD Tagger v3),媒体库展示优化
- 新增 anime_nsfw 模块:WD Tagger v3 模型检测动漫/二次元图片 - 修复误判:当 general 评分最高时判定为安全,避免非动漫图误判 - 数据库存储四个 rating 原始分数(JSON),无需重扫即可查看 - 媒体库新增动漫色情检测列、文件路径列(支持一键复制) - 修复预览图加载占位符显示两个点的问题 - 移动文件按钮不再限制扫描完成后才显示
This commit is contained in:
@@ -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
|
||||
}
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user