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

- 新增 anime_nsfw 模块：WD Tagger v3 模型检测动漫/二次元图片 - 修复误判：当 general 评分最高时判定为安全，避免非动漫图误判 - 数据库存储四个 rating 原始分数(JSON)，无需重扫即可查看 - 媒体库新增动漫色情检测列、文件路径列(支持一键复制) - 修复预览图加载占位符显示两个点的问题 - 移动文件按钮不再限制扫描完成后才显示
2025-08-24 21:14:04 +08:00
parent bd84bc98e8
commit d6798ef1ee
12 changed files with 934 additions and 31 deletions
@@ -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
+}