feat: 色情图片识别(NudeNet)；明暗档位改为5档(20%~100%)

2025-06-02 21:04:38 +08:00
parent ea88f3e766
commit a914b4be74
13 changed files with 450 additions and 25 deletions
@@ -0,0 +1,258 @@
+// Package nsfw 负责"色情图片识别"：本机 Python + NudeNet 起一个本地 HTTP 服务，
+// Go 侧把图片路径发给它，拿到"是否色情"和最高敏感分。
+//
+// 大白话理解：NudeNet 是开源里口碑不错的离线鉴黄模型（python 库，自带模型）。
+// 我们不想让 Go 里掺一堆深度学习依赖，所以让 python 当"外援"：
+//   - 扫描开始时拉起一个常驻的 python 小服务（模型只加载一次）；
+//   - 探测图片时逐张把路径 POST 给它；
+//   - 扫描结束把服务停掉、清理临时脚本。
+//
+// 机器上没有 python 或没装 nudenet 时，Start 会报错，调用方跳过识别、
+// 走默认整理流程——这个功能是纯附加的，不影响其它步骤。
+package nsfw
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"sync"
+	"time"
+
+	_ "embed"
+)
+
+// 内嵌的 python 服务脚本：用标准库起 HTTP 服务，负责加载模型、接收路径、返回结果。
+// 不依赖 flask/fastapi，省去额外装包。
+//
+//go:embed server.py
+var serverScript []byte
+
+// 敏感类集合：检测到这些部位"暴露"且分数达标才算色情。
+// 默认只取最明确的四类（生殖器/肛/女性乳房），避免把泳装、露背照误判。
+var sensitiveClasses = map[string]bool{
+	"FEMALE_GENITALIA_EXPOSED": true,
+	"MALE_GENITALIA_EXPOSED":   true,
+	"ANUS_EXPOSED":             true,
+	"FEMALE_BREAST_EXPOSED":    true,
+}
+
+// detectorStartTimeout 等待 python 服务就绪（模型加载）的超时时间。
+// 首次加载模型可能要十几秒，给足余量。
+const detectorStartTimeout = 60 * time.Second
+
+// requestTimeout 单张图片检测的超时时间。检测很快（毫秒级），
+// 万一卡住（模型没加载好/图太大）不至于拖死整个扫描。
+const requestTimeout = 30 * time.Second
+
+// Detector 封装一个本地 python 服务的客户端。
+// 多个探测协程会并发调 Detect，底层 HTTP 客户端线程安全；
+// 服务端单线程处理，请求自己排队，不用额外加锁。
+type Detector struct {
+	url       string  // 本地服务地址，如 http://127.0.0.1:54321
+	threshold float64 // 判定阈值 0~1
+	client    *http.Client
+	cmd       *exec.Cmd  // python 进程句柄（Stop 时杀掉）
+	scriptDir string     // 临时脚本所在目录（Stop 时清理）
+	mu        sync.Mutex // 保护下面的关闭状态
+	stopped   bool       // 已经 Stop 过了
+	cancel    context.CancelFunc
+	done      chan struct{} // 进程退出时关闭（Wait 协程负责）
+}
+
+// Start 拉起本地 python 检测服务并等它就绪。
+//   - pythonPath：python 可执行文件路径；为空时在 PATH 里找 python/python3。
+//   - threshold：判定阈值 0~1。
+//
+// 找不到 python、装不了 nudenet、或服务起不来时返回错误（调用方跳过识别即可）。
+func Start(pythonPath string, threshold float64) (*Detector, error) {
+	if threshold <= 0 || threshold >= 1 {
+		threshold = 0.5 // 阈值传得不合理就用默认
+	}
+	// 找 python
+	py := resolvePython(pythonPath)
+	if py == "" {
+		return nil, fmt.Errorf("未找到 python，敏感识别不可用")
+	}
+
+	// 先把内嵌的 server.py 写到临时目录
+	tmpDir, err := os.MkdirTemp("", "kmp-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)
+	}
+
+	// 找一个空闲端口
+	port, err := freePort()
+	if err != nil {
+		_ = os.RemoveAll(tmpDir)
+		return nil, fmt.Errorf("找空闲端口失败: %v", err)
+	}
+
+	// 拉起服务。标准输出/错误接到程序自身的输出，方便看 python 报错。
+	ctx, cancel := context.WithCancel(context.Background())
+	cmd := exec.CommandContext(ctx, py, scriptPath, fmt.Sprintf("%d", port))
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+	if err := cmd.Start(); err != nil {
+		cancel()
+		_ = os.RemoveAll(tmpDir)
+		return nil, fmt.Errorf("启动检测服务失败: %v", err)
+	}
+
+	done := make(chan struct{})
+	d := &Detector{
+		url:       fmt.Sprintf("http://127.0.0.1:%d", port),
+		threshold: threshold,
+		client:    &http.Client{Timeout: requestTimeout},
+		cmd:       cmd,
+		scriptDir: tmpDir,
+		cancel:    cancel,
+		done:      done,
+	}
+	// 后台等进程退出：ProcessState 要 Wait 之后才有值，这里用 done 通知，
+	// 这样 waitReady 能及时发现"python 没装 nudenet 直接退掉"的情况。
+	go func() {
+		_ = cmd.Wait()
+		close(done)
+	}()
+
+	// 等服务就绪（模型加载完成）。就绪前先探测 /health。
+	if err := d.waitReady(); err != nil {
+		d.Stop()
+		return nil, err
+	}
+	return d, nil
+}
+
+// waitReady 轮询本地服务，直到能正常应答。
+func (d *Detector) waitReady() error {
+	deadline := time.Now().Add(detectorStartTimeout)
+	for time.Now().Before(deadline) {
+		// 服务进程提前退了（多半是 python 没装 nudenet）直接报错
+		select {
+		case <-d.done:
+			return fmt.Errorf("检测服务提前退出，请确认已安装 nudenet (pip install nudenet)")
+		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)
+}
+
+// Detect 检测一张图片是否色情。返回：
+//   - score：命中的敏感类里最高分数（没有命中为 0）
+//   - nsfw：是否判定为色情（有敏感类且分数 >= 阈值）
+//   - err：检测失败（服务挂了/图读不出等），由调用方兜底按未命中处理
+func (d *Detector) Detect(path string) (score float64, nsfw bool, 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, err
+	}
+	defer resp.Body.Close()
+	data, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return 0, false, err
+	}
+	var out struct {
+		Ok    bool    `json:"ok"`
+		Error string  `json:"error"`
+		Nsfw  bool    `json:"nsfw"`
+		Score float64 `json:"score"`
+	}
+	if err := json.Unmarshal(data, &out); err != nil {
+		return 0, false, fmt.Errorf("解析检测结果失败: %v", err)
+	}
+	if !out.Ok {
+		return 0, false, fmt.Errorf("检测服务返回错误: %s", out.Error)
+	}
+	return out.Score, out.Nsfw, nil
+}
+
+// Stop 停掉检测服务并清理临时文件。可以安全地重复调用。
+func (d *Detector) Stop() {
+	d.mu.Lock()
+	if d.stopped {
+		d.mu.Unlock()
+		return
+	}
+	d.stopped = true
+	d.mu.Unlock()
+
+	if d.cancel != nil {
+		d.cancel()
+	}
+	if d.cmd != nil && d.cmd.Process != nil {
+		// 进程活着才 kill；kill 失败也无所谓，服务很快就没了
+		_ = d.cmd.Process.Kill()
+	}
+	if d.done != nil {
+		// 等 Wait 协程回收进程，避免僵尸进程；进程已死则立刻返回
+		<-d.done
+	}
+	if d.scriptDir != "" {
+		_ = os.RemoveAll(d.scriptDir)
+	}
+}
+
+// resolvePython 在配置值 / PATH 里找可用的 python 解释器。
+// Windows 上常见 python / python3 / py；Linux/mac 常见 python3 / python。
+func resolvePython(configured string) string {
+	candidates := []string{}
+	if strings.TrimSpace(configured) != "" {
+		candidates = append(candidates, strings.TrimSpace(configured))
+	}
+	if os.Getenv("KM_NSFW_PYTHON") != "" {
+		candidates = append(candidates, os.Getenv("KM_NSFW_PYTHON"))
+	}
+	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 ""
+}
+
+// freePort 找一个当前空闲的 TCP 端口（挑到立刻释放，竞态可接受）。
+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
+}