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
+2
View File
@@ -1,3 +1,5 @@
.idea/
.qwen/
media-info.db
__pycache__/
lib/anime_nsfw/
+10 -3
View File
@@ -17,8 +17,15 @@ organize.workers=8
# 模型:优先加载程序运行目录下 lib/nwfs/320n.onnx(推荐随程序分发,离线可用);
# 本地没有会自动下载到该位置,下载失败再回退用 nudenet 自带模型
# NSFW是用来识别真人图片的
nsfw.enabled=true
nsfw.enabled=false
# python 可执行文件路径;留空自动在 PATH 里找 python/python3/py
nsfw.python=C:\\Users\\Administrator\\AppData\\Local\\Programs\\Python\\Python312\\python.exe
# 判定阈值 0~1,敏感部位检测分高于此才算色情(默认 0.2,实测)
nsfw.threshold=0.35
# 判定阈值 0~1,敏感部位检测分高于此才算色情
nsfw.threshold=0.27
# 动漫色情图片识别(需本机安装 python 与 onnxruntime/opencv/numpy:pip install onnxruntime opencv-python numpy)
# true=开启,判定为色情的动漫图片同样移入「图片_敏感_*」目录;false=关闭(默认)
# 模型:WD Tagger v3 (eva02-large),首次启动自动从 HuggingFace 下载到 lib/anime_nsfw/(约 1.26GB)
# AnimeNSFW是用来识别动漫/二次元图片的
anime_nsfw.enabled=true
# 判定阈值 0~1,explicit/questionable 分数高于此才算色情
anime_nsfw.threshold=0.5
+108
View File
@@ -1,7 +1,10 @@
<template>
<div class="page-card">
<div class="toolbar">
<div class="toolbar-left">
<span class="title">媒体库</span>
<span v-if="manualNsfwAvg > 0" class="nsfw-avg">手动标记 NSFW 均值:{{ (manualNsfwAvg * 100).toFixed(2) }}%</span>
</div>
<el-button type="primary" size="small" icon="el-icon-refresh" @click="loadList">刷新</el-button>
</div>
@@ -43,6 +46,9 @@
:preview-src-list="[getPreviewUrl(scope.row.filePath)]"
style="width: 60px; height: 60px"
fit="cover">
<div slot="placeholder" class="image-placeholder">
<i class="el-icon-loading"></i>
</div>
<div slot="error" class="image-error">
<i class="el-icon-picture-outline"></i>
</div>
@@ -55,6 +61,16 @@
<span class="file-name" :title="scope.row.filePath">{{ scope.row.fileName }}</span>
</template>
</el-table-column>
<el-table-column label="路径" min-width="160">
<template slot-scope="scope">
<div class="path-cell">
<span class="file-path" :title="scope.row.filePath">{{ scope.row.filePath }}</span>
<el-tooltip content="复制路径" placement="top">
<i class="el-icon-document-copy copy-icon" @click="copyPath(scope.row.filePath)"/>
</el-tooltip>
</div>
</template>
</el-table-column>
<el-table-column label="类型" width="80" align="center">
<template slot-scope="scope">
<el-tag :type="typeTag(scope.row.mediaType)" size="small">{{ typeName(scope.row.mediaType) }}</el-tag>
@@ -89,6 +105,21 @@
<span v-else class="muted">-</span>
</template>
</el-table-column>
<el-table-column label="动漫色情检测" width="140" align="center">
<template slot-scope="scope">
<div v-if="scope.row.animeNsfwScore > 0">
<el-tag :type="scope.row.animeNsfwScore >= scope.row.animeNsfwThreshold ? 'danger' : 'warning'" size="small">
{{ (scope.row.animeNsfwScore * 100).toFixed(1) }}%
</el-tag>
<div class="nsfw-threshold">阈值: {{ (scope.row.animeNsfwThreshold * 100).toFixed(0) }}%</div>
</div>
<div v-else-if="scope.row.animeNsfwRatings">
<el-tag type="info" size="small">安全</el-tag>
<div class="nsfw-threshold">{{ formatAnimeRatings(scope.row.animeNsfwRatings) }}</div>
</div>
<span v-else class="muted">-</span>
</template>
</el-table-column>
<el-table-column label="色情标记" width="100" align="center">
<template slot-scope="scope">
<el-switch
@@ -159,6 +190,7 @@ export default {
data() {
return {
loading: false,
manualNsfwAvg: 0,
list: [],
total: 0,
resolutionOptions: RES_STANDARD.concat(resNonStandard()),
@@ -182,6 +214,7 @@ export default {
mounted() {
this.loadOptions()
this.loadList()
this.loadManualNsfwAvg()
},
methods: {
async loadOptions() {
@@ -255,15 +288,52 @@ export default {
try {
await request.put(`/api/media/${row.id}/manual-nsfw`, { value: newValue })
row.manualNsfw = newValue
this.loadManualNsfwAvg()
} catch (e) {
// 拦截器已提示
}
},
async loadManualNsfwAvg() {
try {
const res = await request.get('/api/media/manual-nsfw-avg')
this.manualNsfwAvg = (res.data && res.data.avg) || 0
} catch (e) {
// 忽略
}
},
copyPath(path) {
navigator.clipboard.writeText(path).then(() => {
this.$message.success('路径已复制')
}).catch(() => {
this.$message.error('复制失败')
})
},
formatAnimeRatings(ratingsJson) {
if (!ratingsJson) return ''
try {
const r = JSON.parse(ratingsJson)
return `G=${(r.general * 100).toFixed(0)} S=${(r.sensitive * 100).toFixed(0)} Q=${(r.questionable * 100).toFixed(0)} E=${(r.explicit * 100).toFixed(0)}`
} catch (e) {
return ''
}
}
}
}
</script>
<style scoped>
.toolbar-left {
display: flex;
align-items: center;
gap: 12px;
}
.nsfw-avg {
font-size: 13px;
color: #F56C6C;
font-weight: 500;
}
.filter-bar {
display: flex;
gap: 10px;
@@ -276,6 +346,44 @@ export default {
color: #1f3a6b;
}
.file-path {
color: #606266;
font-size: 12px;
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
flex: 1;
}
.path-cell {
display: flex;
align-items: center;
gap: 4px;
overflow: hidden;
}
.copy-icon {
cursor: pointer;
color: #909399;
font-size: 14px;
flex-shrink: 0;
}
.copy-icon:hover {
color: #409EFF;
}
.image-placeholder {
display: flex;
align-items: center;
justify-content: center;
width: 60px;
height: 60px;
background: #f5f7fa;
color: #c0c4cc;
font-size: 20px;
}
.muted {
color: #c0c4cc;
}
+1 -2
View File
@@ -58,13 +58,12 @@
:loading="starting"
@click="handleStart">开始扫描</el-button>
<el-button
v-if="status.state === 'done'"
type="warning"
size="medium"
icon="el-icon-s-promotion"
:loading="moving"
@click="handleMove">移动文件到存放目录</el-button>
<div v-else-if="status.state === 'running'" class="running-hint">
<div v-if="status.state === 'running'" class="running-hint">
<span>
<i class="el-icon-loading"></i>
<template v-if="status.stage === 'move'">
+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()
+10
View File
@@ -137,6 +137,16 @@ func (c *MediaController) Preview(ctx *gin.Context) {
ctx.File(path)
}
// AvgManualNsfwScore 处理 GET /api/media/manual-nsfw-avg:返回手动标记均值。
func (c *MediaController) AvgManualNsfwScore(ctx *gin.Context) {
avg, err := c.store.AvgManualNsfwScore()
if err != nil {
Error(ctx, http.StatusInternalServerError, "查询失败:"+err.Error())
return
}
Success(ctx, gin.H{"avg": avg})
}
// UpdateManualNsfw 处理 PUT /api/media/:id/manual-nsfw:更新手动色情标记。
func (c *MediaController) UpdateManualNsfw(ctx *gin.Context) {
idStr := ctx.Param("id")
+1
View File
@@ -34,6 +34,7 @@ func NewRouter(st *store.Store) *gin.Engine {
apiGroup.POST("/media/stats", mc.Stats)
apiGroup.GET("/media/preview", mc.Preview)
apiGroup.PUT("/media/:id/manual-nsfw", mc.UpdateManualNsfw)
apiGroup.GET("/media/manual-nsfw-avg", mc.AvgManualNsfwScore)
// 扫描整理(开始 / 中止 / 移动 / 查进度 / 实时进度)
sc := NewScanController(service.NewScanService(st))
+12 -1
View File
@@ -22,6 +22,8 @@ import (
// - NSFWEnabled:是否开启色情图片识别(本地 python + NudeNet,默认 false)
// - NSFWPython:python 可执行文件路径;留空自动在 PATH 里找
// - NSFWThreshold:判定阈值 0~1(默认 0.5)
// - AnimeNSFWEnabled:是否开启动漫色情图片识别(本地 python + WD Tagger v3,默认 false)
// - AnimeNSFWThreshold:动漫判定阈值 0~1(默认 0.5)
type Config struct {
Port int
Dsn string
@@ -32,6 +34,8 @@ type Config struct {
NSFWEnabled bool
NSFWPython string
NSFWThreshold float64
AnimeNSFWEnabled bool
AnimeNSFWThreshold float64
}
// appConfig 当前生效的配置(程序运行期间一直不变)。
@@ -45,7 +49,7 @@ func GetAppConfig() Config {
// LoadAppConfig 从 application.properties 读配置;文件不存在就用默认值。
func LoadAppConfig() {
// 先给一套默认值
appConfig = Config{Port: 8080, Dsn: "media-info.db", AutoOpen: true, RotateImages: true, Workers: 4, NSFWThreshold: 0.5}
appConfig = Config{Port: 8080, Dsn: "media-info.db", AutoOpen: true, RotateImages: true, Workers: 4, NSFWThreshold: 0.5, AnimeNSFWThreshold: 0.5}
p, err := properties.LoadFile("application.properties", properties.UTF8)
if err != nil {
@@ -80,6 +84,13 @@ func LoadAppConfig() {
appConfig.NSFWThreshold = t
}
}
// 动漫色情识别:默认关;阈值只认 0~1 之间的小数
appConfig.AnimeNSFWEnabled = p.GetBool("anime_nsfw.enabled", false)
if v, ok := p.Get("anime_nsfw.threshold"); ok && strings.TrimSpace(v) != "" {
if t, err := strconv.ParseFloat(strings.TrimSpace(v), 64); err == nil && t > 0 && t < 1 {
appConfig.AnimeNSFWThreshold = t
}
}
}
// WorkerCount 返回探测/MIME/移动阶段使用的并发 worker 数:
+60 -4
View File
@@ -12,6 +12,7 @@ package organize
import (
"context"
"encoding/json"
"fmt"
"os"
"path/filepath"
@@ -20,6 +21,7 @@ import (
"sync/atomic"
"time"
"kenaito-media-picker/internal/anime_nsfw"
"kenaito-media-picker/internal/classify"
"kenaito-media-picker/internal/config"
"kenaito-media-picker/internal/digest"
@@ -71,6 +73,7 @@ type Organizer struct {
mover *mover.Mover // 负责实际搬文件的家伙
ctx context.Context // 中止信号;ctx 被取消时探测/归类立即停下
nsfw *nsfw.Detector // 色情检测器;为 nil 表示没启用,图片不检测
animeNsfw *anime_nsfw.Detector // 动漫色情检测器;为 nil 表示没启用
rotateImages bool // 是否处理 EXIF 旋转(转正/按显示方向归类);false 则完全不动,按存储方向归类
buckets map[string][]string // 目标目录 → 里面该放哪些文件(只记源路径)
order []string // 目标目录的先后顺序(保证输出稳定、不来回跳)
@@ -81,11 +84,14 @@ type Organizer struct {
pendingScan []store.ScanLog // 待写入续传表的记录(攒一批再批量写)
nsfwChecked int32 // 色情识别已检查的图片数(并行累计,只涨不跌)
nsfwHits int32 // 色情识别命中的图片数(并行累计)
animeNsfwChecked int32 // 动漫色情识别已检查的图片数
animeNsfwHits int32 // 动漫色情识别命中的图片数
// 阶段耗时统计(并行 worker 累计,单位纳秒,只用于日志诊断瓶颈)
tImageProbe int64 // 图片尺寸探测总耗时
tImgMd5 int64 // 图片 MD5/指纹计算总耗时
tBright int64 // 明暗分析总耗时
tNsfw int64 // 色情识别总耗时
tAnimeNsfw int64 // 动漫色情识别总耗时
tMediaProbe int64 // 媒体探测(ffprobe)总耗时
tMediaMd5 int64 // 媒体(视频/音频) MD5/指纹计算总耗时
}
@@ -128,6 +134,12 @@ func (o *Organizer) SetNSFWDetector(d *nsfw.Detector) {
o.nsfw = d
}
// SetAnimeNSFWDetector 设置动漫色情检测器。传 nil 表示不检测(默认)。
// 设置了检测器后,图片探测时会额外做动漫色情识别,命中的图片同样归入「图片_敏感_*」目录。
func (o *Organizer) SetAnimeNSFWDetector(d *anime_nsfw.Detector) {
o.animeNsfw = d
}
// SetRotateImages 设置是否处理图片 EXIF 旋转(默认 true)。
// - true:带旋转标记的图片做转正处理(有 ffmpeg 就烧进像素,没有则按显示方向归类);
// - false:完全不动,图片按文件里存储的方向/尺寸归类,不做任何旋转相关处理。
@@ -394,6 +406,9 @@ type imageProbeResult struct {
lightScore float64 // 亮色像素占比
nsfwScore float64 // 色情检测最高分(未检测为 0)
nsfw bool // 是否判定为色情(检测可用时才可能为 true)
animeNsfwScore float64 // 动漫色情检测最高分(未检测为 0)
animeNsfw bool // 是否判定为动漫色情
animeNsfwRatings anime_nsfw.AnimeRatings // WD Tagger 四个 rating 的原始分数
}
// probeImage 给一张图片做体检(并行阶段执行):探测尺寸 + 算 MD5 + 分析明暗。
@@ -470,6 +485,29 @@ func (o *Organizer) probeImage(f scanner.File) imageProbeResult {
}
}
}
// 动漫色情识别(GIF 不做;没启用检测器就直接跳过)
if o.animeNsfw != nil && !strings.EqualFold(f.Ext, ".gif") {
animeStart := time.Now()
score, hit, ratings, derr := o.animeNsfw.Detect(f.Path)
atomic.AddInt64(&o.tAnimeNsfw, int64(time.Since(animeStart)))
if derr != nil {
fmt.Printf("=== 动漫色情检测失败(按未命中归类): %s, 错误: %v \n", f.Path, derr)
} else {
r.animeNsfwScore = score
r.animeNsfw = hit
r.animeNsfwRatings = ratings
checked := atomic.AddInt32(&o.animeNsfwChecked, 1)
if hit {
atomic.AddInt32(&o.animeNsfwHits, 1)
fmt.Printf("=== 动漫色情命中: %s, 检测分=%.4f (阈值=%.4f)\n", f.Path, score, o.animeNsfw.Threshold())
} else {
fmt.Printf("=== 动漫色情未命中: %s, 检测分=%.4f (阈值=%.4f)\n", f.Path, score, o.animeNsfw.Threshold())
}
if checked%nsfwProgressInterval == 0 {
fmt.Printf("=== 动漫色情识别进度: 已检查 %d 张, 命中 %d 张\n", checked, atomic.LoadInt32(&o.animeNsfwHits))
}
}
}
return r
}
@@ -520,6 +558,9 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
if o.nsfw != nil {
fmt.Printf("=== 色情识别: 共 %d 张图片待检测(GIF跳过), 每 %d 张报一次进度\n", total, nsfwProgressInterval)
}
if o.animeNsfw != nil {
fmt.Printf("=== 动漫色情识别: 共 %d 张图片待检测(GIF跳过), 每 %d 张报一次进度\n", total, nsfwProgressInterval)
}
results := runParallel(o.ctx, images, o.probeImage, func(done int) {
o.report(PhaseImages, done, total)
})
@@ -561,9 +602,19 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
info.NsfwScore = r.nsfwScore
info.NsfwThreshold = o.nsfw.Threshold()
}
if r.nsfw {
if info != nil && r.animeNsfwScore > 0 {
info.AnimeNsfwScore = r.animeNsfwScore
info.AnimeNsfwThreshold = o.animeNsfw.Threshold()
if ratingsJSON, err := json.Marshal(r.animeNsfwRatings); err == nil {
info.AnimeNsfwRatings = string(ratingsJSON)
}
}
// 真人色情或动漫色情任一命中,都归入敏感目录
if r.nsfw || r.animeNsfw {
o.addClassified(f, classify.ImageDirSensitive(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), r.size, r.headMd5, r.tailMd5)
if r.nsfw {
nsfwCount++
}
} else {
o.addClassified(f, classify.ImageDir(f.Ext, r.res.Width, r.res.Height, r.darkScore, r.lightScore), r.size, r.headMd5, r.tailMd5)
}
@@ -576,14 +627,19 @@ func (o *Organizer) ProcessImages(images []scanner.File) {
fmt.Printf("=== 续传跳过: %d 个文件(上次已归类,直接复用)\n", skip)
}
fmt.Printf("=== 图片处理完毕,成功: %d, 读取异常: %d, 忽略: %d, 重复: %d \n\n", success, fail, ignore, dup)
fmt.Printf("=== 图片归类耗时: 实际 %.1fs, 其中(并行累计) 探测 %.1fs, MD5 %.1fs, 明暗 %.1fs, 色情 %.1fs\n",
fmt.Printf("=== 图片归类耗时: 实际 %.1fs, 其中(并行累计) 探测 %.1fs, MD5 %.1fs, 明暗 %.1fs, 色情 %.1fs, 动漫色情 %.1fs\n",
time.Since(phaseStart).Seconds(),
time.Duration(atomic.LoadInt64(&o.tImageProbe)).Seconds(),
time.Duration(atomic.LoadInt64(&o.tImgMd5)).Seconds(),
time.Duration(atomic.LoadInt64(&o.tBright)).Seconds(),
time.Duration(atomic.LoadInt64(&o.tNsfw)).Seconds())
time.Duration(atomic.LoadInt64(&o.tNsfw)).Seconds(),
time.Duration(atomic.LoadInt64(&o.tAnimeNsfw)).Seconds())
animeNsfwHits := atomic.LoadInt32(&o.animeNsfwHits)
if o.nsfw != nil && nsfwCount > 0 {
fmt.Printf("=== 色情识别: 共命中 %d 张,已归入「图片_敏感_*」目录,可在【色情识别】页面查看并人工复核\n", nsfwCount)
fmt.Printf("=== 色情识别: 共命中 %d 张,已归入「图片_敏感_*」目录,可在【媒体库】页面查看并通过<色情标记>人工复核\n", nsfwCount)
}
if o.animeNsfw != nil && animeNsfwHits > 0 {
fmt.Printf("=== 动漫色情识别: 共命中 %d 张,已归入「图片_敏感_*」目录,可在【媒体库】页面查看并通过<色情标记>人工复核\n", animeNsfwHits)
}
}
+21 -2
View File
@@ -23,6 +23,7 @@ import (
"sync"
"time"
"kenaito-media-picker/internal/anime_nsfw"
"kenaito-media-picker/internal/config"
"kenaito-media-picker/internal/mover"
"kenaito-media-picker/internal/nsfw"
@@ -249,7 +250,7 @@ func (s *ScanService) Cancel() error {
// Move 触发文件移动操作。
// 根据 scan_log 中记录的分类结果,将文件从源路径移动到目标目录。
// 只有在扫描完成(state == done)且当前没有任务运行时才能调用。
// 只有在扫描完成(state == done)或 scan_log 有记录时才能调用。
func (s *ScanService) Move() error {
s.mu.Lock()
@@ -259,11 +260,15 @@ func (s *ScanService) Move() error {
return errors.New("扫描整理正在运行中,请等待完成")
}
// 必须先完成一次扫描,才有 scan_log 记录
// 必须先完成一次扫描,或者 scan_log 有记录,才有数据可移动
if s.state != ScanDone {
// 检查 scan_log 是否有记录
hasLog, err := s.store.HasScanLog()
if err != nil || !hasLog {
s.mu.Unlock()
return errors.New("请先完成扫描后再执行移动")
}
}
// 把状态改成"运行中",清空上一次的残留信息
s.state = ScanRunning
@@ -681,8 +686,22 @@ func (s *ScanService) runScan(ctx context.Context, rootDir, destDir string) erro
}
}
// ---------- 第六步(可选):动漫色情识别 ----------
// 用本机 python + WD Tagger v3 起一个本地检测服务。没装 python/onnxruntime 时跳过。
animeNsfwDet := (*anime_nsfw.Detector)(nil)
if cfg := config.GetAppConfig(); cfg.AnimeNSFWEnabled {
if d, err := anime_nsfw.Start(cfg.NSFWPython, cfg.AnimeNSFWThreshold); err != nil {
fmt.Printf("=== 动漫色情识别未启用(图片照常处理): %v\n", err)
} else {
animeNsfwDet = d
defer animeNsfwDet.Stop()
fmt.Println("=== 动漫色情识别已启用(WD Tagger v3)")
}
}
org := organize.New(pickerDir, ffprobePath, ffmpegPath, s.store, ctx)
org.SetNSFWDetector(nsfwDet)
org.SetAnimeNSFWDetector(animeNsfwDet)
org.SetRotateImages(config.GetAppConfig().RotateImages)
// 整个"整理"分两个阶段,界面上共用一条进度条:
+36
View File
@@ -49,6 +49,9 @@ type MediaInfo struct {
TailMd5 string `xorm:"char(32) index('idx_fp') 'tail_md5'" json:"tailMd5"` // 文件尾 256KB 的 MD5(指纹用)
NsfwScore float64 `xorm:"real 'nsfw_score'" json:"nsfwScore"` // NSFW 检测分 0~1(0 表示未检测或非图片)
NsfwThreshold float64 `xorm:"real 'nsfw_threshold'" json:"nsfwThreshold"` // NSFW 判定阈值
AnimeNsfwScore float64 `xorm:"real 'anime_nsfw_score'" json:"animeNsfwScore"` // 动漫 NSFW 检测分 0~1(0 表示未检测或非动漫)
AnimeNsfwThreshold float64 `xorm:"real 'anime_nsfw_threshold'" json:"animeNsfwThreshold"` // 动漫 NSFW 判定阈值
AnimeNsfwRatings string `xorm:"text 'anime_nsfw_ratings'" json:"animeNsfwRatings"` // WD Tagger 四个 rating 的 JSON(G/S/Q/E)
ManualNsfw string `xorm:"varchar(1) 'manual_nsfw'" json:"manualNsfw"` // 手动标记是否色情:Y/N,默认 N
CreatedAt time.Time `xorm:"created 'created_at'" json:"createdAt"` // 首次入库时间
UpdatedAt time.Time `xorm:"updated 'updated_at'" json:"updatedAt"` // 最近更新时间
@@ -190,6 +193,18 @@ func (s *Store) SaveScanLogs(entries []ScanLog) error {
return sess.Commit()
}
// HasScanLog 检查是否有扫描记录(用于判断是否可以直接移动文件)。
func (s *Store) HasScanLog() (bool, error) {
if s == nil || s.engine == nil {
return false, nil
}
total, err := s.engine.Count(new(ScanLog))
if err != nil {
return false, err
}
return total > 0, nil
}
// ClearScanLog 清空全部续传记录。移动完成后调用。
func (s *Store) ClearScanLog() error {
if s == nil || s.engine == nil {
@@ -220,6 +235,27 @@ func (s *Store) UpdateMediaFilePaths(moved map[string]string) error {
return sess.Commit()
}
// AvgManualNsfwScore 计算 manual_nsfw='Y' 的记录的 nsfw_score 均值。
func (s *Store) AvgManualNsfwScore() (float64, error) {
if s == nil || s.engine == nil {
return 0, nil
}
var result struct {
Avg float64 `xorm:"avg"`
}
has, err := s.engine.Table("media_info").
Select("AVG(nsfw_score) as avg").
Where("manual_nsfw = 'Y'").
Get(&result)
if err != nil {
return 0, err
}
if !has {
return 0, nil
}
return result.Avg, nil
}
// UpdateManualNsfw 更新单条记录的手动色情标记。
func (s *Store) UpdateManualNsfw(id int64, value string) error {
if s == nil || s.engine == nil {