kenaito-media-picker/internal/nsfw/detector.go

// Package nsfw 负责"色情图片识别"：本机 Python + NudeNet 起一个本地 HTTP 服务，
// Go 侧把图片路径发给它，拿到"是否色情"和最高敏感分。
//
// 大白话理解：NudeNet 是开源里口碑不错的离线鉴黄模型（python 库，自带模型）。
// 我们不想让 Go 里掺一堆深度学习依赖，所以让 python 当"外援"：
//   - 扫描开始时拉起一个常驻的 python 小服务（模型只加载一次）；
//   - 探测图片时逐张把路径 POST 给它；
//   - 扫描结束把服务停掉、清理临时脚本。
//
// 机器上没有 python 或没装 nudenet 时，Start 会报错，调用方跳过识别、
// 走默认整理流程——这个功能是纯附加的，不影响其它步骤。
package nsfw

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"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

// 注意：哪些敏感类才算色情，定义在 server.py 的 SENSITIVE 集合里（单一事实来源），
// Go 侧只负责把路径发给服务、读取判定结果，不重复维护这份清单。

// NsfwDetection NudeNet 检测到的一个目标（类别+置信度+位置）。
type NsfwDetection struct {
	Class string  `json:"class"`
	Score float64 `json:"score"`
	Box   [4]int  `json:"box"` // [x, y, w, h]
}

// detectorStartTimeout 等待 python 服务就绪（模型加载）的超时时间。
// 首次启动若本地没有模型还要联网下载约 11.6MB，给足余量。
const detectorStartTimeout = 180 * time.Second

// requestTimeout 单张图片检测的超时时间。检测很快（毫秒级），
// 万一卡住（模型没加载好/图太大）不至于拖死整个扫描。
const requestTimeout = 30 * time.Second

// Detector 封装一个本地 python 服务的客户端。
// 多个探测协程会并发调 Detect，底层 HTTP 客户端线程安全；
// 服务端用固定大小线程池并行推理（每个 worker 一份独立模型实例），请求不必排队。
type Detector struct {
	url       string // 本地服务地址，如 http://127.0.0.1:54321
	workers   int    // Python 推理 worker 数量
	client    *http.Client
	cmd       *exec.Cmd  // python 进程句柄（Stop 时杀掉）
	scriptDir string     // 临时脚本所在目录（Stop 时清理）
	python    string     // python 可执行文件路径（服务重启时要用）
	modelDir  string     // 模型目录 lib（服务重启时要用）
	mu        sync.Mutex // 保护 stopped / 进程句柄
	stopped   bool       // 已经 Stop 过了
	cancel    context.CancelFunc
	done      chan struct{} // 进程退出时关闭（Wait 协程负责）
	restartMu sync.Mutex    // 串行化重启：同一时刻只一个 goroutine 重启，其余排队等
}

// Start 拉起本地 python 检测服务并等它就绪。
//   - pythonPath：python 可执行文件路径；为空时在 PATH 里找 python/python3。
//   - workers：Python 推理 worker 数量；<=0 表示使用默认值（4）。
//
// 找不到 python、装不了 nudenet、或服务起不来时返回错误（调用方跳过识别即可）。
func Start(pythonPath string, workers int) (*Detector, error) {
	if workers <= 0 {
		workers = 4 // 默认 4 个 worker
	}
	// 找 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)
	}

	// 模型目录：程序运行目录下的 lib（和 application.properties / media-info.db 同一根目录）。
	// 服务脚本优先加载 lib/human_nwfs/320n.onnx；本地没有会自己下载到那里，下载失败再回退用 pip 自带模型。
	modelDir := ""
	if wd, err := os.Getwd(); err == nil {
		modelDir = filepath.Join(wd, "lib")
	}

	d := &Detector{
		workers:   workers,
		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
}

// launch 启动（或重启）python 检测服务，并等它就绪（模型加载完成）。
// 调用方需要保证持有 mu，且 stopped 为 false。
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)
	// 设置环境变量，传递 worker 数量给 Python 服务
	cmd.Env = append(os.Environ(), fmt.Sprintf("KM_NSFW_WORKERS=%d", d.workers))
	// 标准输出/错误都接到程序自身的输出：
	//   - 扫描期间 os.Stdout 会被接管进"运行日志"，这样 python 的报错（崩溃原因）也能在前端看到
	//   - 方便排查服务为什么挂掉
	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
	// 后台等进程退出：ProcessState 要 Wait 之后才有值，这里用 done 通知，
	// 这样 waitReady 能及时发现"python 没装 nudenet 直接退掉"的情况。
	go func() {
		_ = cmd.Wait()
		close(done)
	}()

	// 等服务就绪（模型加载完成）。就绪前先探测 /health。
	if err := d.waitReady(); err != nil {
		d.teardown()
		return err
	}
	return nil
}

// teardown 停掉当前 python 进程并清理句柄（不删临时目录）。
// 调用方需要保证持有 mu。进程可能已经死了（崩溃），kill 失败也无所谓。
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 {
		// 等 Wait 协程回收进程，避免僵尸进程；进程已死则立刻返回。
		// 等超过 3 秒就放弃，防止极端情况下卡死重启流程。
		select {
		case <-d.done:
		case <-time.After(3 * time.Second):
		}
		d.done = nil
	}
	d.cmd = 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）
//   - detections：所有检测到的目标（类别+分数+位置）
//   - err：检测失败（服务挂了/图读不出等），由调用方兜底按未命中处理
//
// 服务失联（连接被拒，多半是 python 进程在处理某张图时崩了）时，会串行化地
// 重启服务并重试一次：多个并发 worker 里只有一个负责重启，其余排队，等重启完
// 也重试自己的图——保证服务短暂挂掉不会漏掉一整批图片。
func (d *Detector) Detect(path string) (score float64, detections []NsfwDetection, err error) {
	score, detections, err = d.detectOnce(path)
	if err == nil {
		return score, detections, nil
	}
	// 只有"发请求都失败"（连接被拒/重置，说明服务失联）才走重启；
	// 服务正常返回的明确错误（如某张图读不了）不重启。
	var ce *connError
	if !errors.As(err, &ce) {
		return score, detections, err
	}

	d.restartMu.Lock()
	defer d.restartMu.Unlock()

	if d.isStopped() {
		return 0, nil, err
	}
	// 可能排在前面的 goroutine 已经把服务重启好了：先探活，活着就直接重试，别反复重启
	if d.healthy() {
		return d.detectOnce(path)
	}
	fmt.Printf("=== 色情检测服务失联(%v), 正在重启…\n", ce.err)
	if rerr := d.launch(); rerr != nil {
		return 0, nil, fmt.Errorf("检测服务重启失败: %v", rerr)
	}
	fmt.Println("=== 色情检测服务已重启，继续识别")
	return d.detectOnce(path)
}

// healthy 快速探活：本地服务是否还能正常应答（带短超时，避免拖慢扫描）。
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
}

// isStopped 判断检测器是否已被 Stop（线程安全）。
func (d *Detector) isStopped() bool {
	d.mu.Lock()
	defer d.mu.Unlock()
	return d.stopped
}

// detectOnce 发一次检测请求（不带重启逻辑）。
func (d *Detector) detectOnce(path string) (score float64, detections []NsfwDetection, err error) {
	body, _ := json.Marshal(map[string]interface{}{
		"path": path,
	})
	resp, err := d.client.Post(d.url+"/detect", "application/json", bytes.NewReader(body))
	if err != nil {
		// 打上"连接失败"标记，供 Detect 判断是否要重启
		return 0, nil, &connError{err}
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return 0, nil, err
	}
	var out struct {
		Ok         bool                     `json:"ok"`
		Error      string                   `json:"error"`
		Score      float64                  `json:"score"`
		Detections []map[string]interface{} `json:"detections"`
	}
	if err := json.Unmarshal(data, &out); err != nil {
		return 0, nil, fmt.Errorf("解析检测结果失败: %v", err)
	}
	if !out.Ok {
		return 0, nil, fmt.Errorf("检测服务返回错误: %s", out.Error)
	}
	// 解析检测详情
	detections = make([]NsfwDetection, 0, len(out.Detections))
	for _, det := range out.Detections {
		cls, _ := det["class"].(string)
		sc, _ := det["score"].(float64)
		boxArr, _ := det["box"].([]interface{})
		var box [4]int
		for i := 0; i < 4 && i < len(boxArr); i++ {
			if v, ok := boxArr[i].(float64); ok {
				box[i] = int(v)
			}
		}
		detections = append(detections, NsfwDetection{Class: cls, Score: sc, Box: box})
	}
	// 打印所有检测到的类别和分数（调试用）
	if len(detections) > 0 {
		var classes []string
		for _, det := range detections {
			classes = append(classes, fmt.Sprintf("%s:%.3f", det.Class, det.Score))
		}
		fmt.Printf("=== NSFW 检测详情 %s: [%s]\n", path, strings.Join(classes, ", "))
	} else {
		fmt.Printf("=== NSFW 检测详情 %s: [无检测结果]\n", path)
	}
	return out.Score, detections, nil
}

// connError 标记"请求根本没发出去/连不上本地服务"的错误（服务失联）。
type connError struct{ err error }

func (e *connError) Error() string { return e.err.Error() }
func (e *connError) Unwrap() error { return e.err }

// Stop 停掉检测服务并清理临时文件。可以安全地重复调用。
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)
	}
}

// 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
}