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

// Package hentai 负责"二次元图片识别"：本机 Python + wd-tagger 起一个本地 HTTP 服务，
// Go 侧把图片路径发给它，拿到"是否二次元敏感"和分级信息。
//
// 大白话理解：wd-tagger 是专门针对二次元训练的分级模型，能区分 safe/questionable/explicit。
// 我们让 python 当"外援"：
//   - 扫描开始时拉起一个常驻的 python 小服务（模型只加载一次）；
//   - 探测图片时逐张把路径 POST 给它；
//   - 扫描结束把服务停掉、清理临时脚本。
//
// 机器上没有 python 或没装 onnxruntime 时，Start 会报错，调用方跳过识别、
// 走普通整理流程——这个功能是纯附加的，不影响其它步骤。
package hentai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"sync"
	"time"

	_ "embed"
)

// 内嵌的 python 服务脚本。
//go:embed server.py
var serverScript []byte

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

// requestTimeout 单张图片检测的超时时间。
const requestTimeout = 120 * time.Second

// Detector 封装一个本地 python 服务的客户端。
type Detector struct {
	url       string  // 本地服务地址，如 http://127.0.0.1:54322
	threshold float64 // 判定阈值 0~1
	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 重启，其余排队等
}

// DetectResult 二次元检测的结果。
type DetectResult struct {
	Hentai     bool               // 是否判定为二次元敏感（explicit 或 questionable 超过阈值）
	Score      float64            // 最高非 safe 分数（explicit 或 questionable 中较高的）
	Rating     string             // 分级标签：safe / questionable / explicit
	Ratings    map[string]float64 // 三个分级的原始分数
	TopTags    []TagScore         // 高置信度标签（最多 10 个）
}

// TagScore 一个标签及其分数。
type TagScore struct {
	Tag   string  `json:"tag"`
	Score float64 `json:"score"`
}

// Start 拉起本地 python 检测服务并等它就绪。
//   - pythonPath：python 可执行文件路径；为空时在 PATH 里找 python/python3。
//   - threshold：判定阈值 0~1。
//
// 找不到 python、装不了依赖、或服务起不来时返回错误（调用方跳过识别即可）。
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-hentai-*")
	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 同一根目录）。
	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
}

// launch 启动（或重启）python 检测服务，并等它就绪（模型加载完成）。
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)
	}()

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

// teardown 停掉当前 python 进程并清理句柄（不删临时目录）。
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
}

// waitReady 轮询本地服务，直到能正常应答。
func (d *Detector) waitReady() error {
	deadline := time.Now().Add(detectorStartTimeout)
	for time.Now().Before(deadline) {
		// 服务进程提前退了直接报错
		select {
		case <-d.done:
			return fmt.Errorf("二次元检测服务提前退出，请确认已安装 onnxruntime 和 Pillow (pip install onnxruntime Pillow 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(1 * time.Second) // 模型加载较慢，轮询间隔拉长
	}
	return fmt.Errorf("二次元检测服务在 %v 内未就绪", detectorStartTimeout)
}

// Detect 检测一张图片是否为二次元敏感内容。
// 服务失联时会串行化地重启服务并重试一次。
func (d *Detector) Detect(path string) (DetectResult, error) {
	result, err := d.detectOnce(path)
	if err == nil {
		return result, nil
	}
	// 只有"发请求都失败"才走重启
	var ce *connError
	if !errors.As(err, &ce) {
		return DetectResult{}, err
	}

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

	if d.isStopped() {
		return DetectResult{}, err
	}
	// 可能排在前面的 goroutine 已经把服务重启好了：先探活，活着就直接重试
	if d.healthy() {
		return d.detectOnce(path)
	}
	fmt.Printf("=== 二次元检测服务失联(%v), 正在重启…\n", ce.err)
	if rerr := d.launch(); rerr != nil {
		return DetectResult{}, 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) (DetectResult, 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 DetectResult{}, &connError{err}
	}
	defer resp.Body.Close()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return DetectResult{}, err
	}
	var out struct {
		Ok      bool               `json:"ok"`
		Error   string             `json:"error"`
		Hentai  bool               `json:"hentai"`
		Score   float64            `json:"score"`
		Rating  string             `json:"rating"`
		Ratings map[string]float64 `json:"ratings"`
		TopTags []TagScore         `json:"topTags"`
	}
	if err := json.Unmarshal(data, &out); err != nil {
		return DetectResult{}, fmt.Errorf("解析检测结果失败: %v", err)
	}
	if !out.Ok {
		return DetectResult{}, fmt.Errorf("二次元检测服务返回错误: %s", out.Error)
	}
	return DetectResult{
		Hentai:  out.Hentai,
		Score:   out.Score,
		Rating:  out.Rating,
		Ratings: out.Ratings,
		TopTags: out.TopTags,
	}, 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)
	}
}

// Threshold 返回当前使用的判定阈值（0~1）。给上层记录结果用。
func (d *Detector) Threshold() float64 {
	return d.threshold
}

// resolvePython 在配置值 / PATH 里找可用的 python 解释器。
func resolvePython(configured string) string {
	candidates := []string{}
	if strings.TrimSpace(configured) != "" {
		candidates = append(candidates, strings.TrimSpace(configured))
	}
	if os.Getenv("KM_HENTAI_PYTHON") != "" {
		candidates = append(candidates, os.Getenv("KM_HENTAI_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
}