265 lines
8.5 KiB
Go
265 lines
8.5 KiB
Go
// 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
|
||
|
||
// 注意：哪些敏感类才算色情，定义在 server.py 的 SENSITIVE 集合里（单一事实来源），
|
||
// Go 侧只负责把路径发给服务、读取判定结果，不重复维护这份清单。
|
||
|
||
// detectorStartTimeout 等待 python 服务就绪（模型加载）的超时时间。
|
||
// 首次启动若本地没有模型还要联网下载约 11.6MB，给足余量。
|
||
const detectorStartTimeout = 180 * 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)
|
||
}
|
||
|
||
// 模型目录：程序运行目录下的 lib（和 application.properties / media-info.db 同一根目录）。
|
||
// 服务脚本优先加载 lib/320n.onnx；本地没有会自己下载到那里，下载失败再回退用 pip 自带模型。
|
||
modelDir := ""
|
||
if wd, err := os.Getwd(); err == nil {
|
||
modelDir = filepath.Join(wd, "lib")
|
||
}
|
||
|
||
// 拉起服务。标准输出/错误接到程序自身的输出，方便看 python 报错。
|
||
ctx, cancel := context.WithCancel(context.Background())
|
||
cmd := exec.CommandContext(ctx, py, scriptPath, fmt.Sprintf("%d", port), modelDir)
|
||
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)
|
||
}
|
||
}
|
||
|
||
// Threshold 返回当前使用的判定阈值（0~1）。给上层记录结果用。
|
||
func (d *Detector) Threshold() float64 {
|
||
return d.threshold
|
||
}
|
||
|
||
// 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
|
||
}
|