342 lines
8.7 KiB
Go
342 lines
8.7 KiB
Go
// 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
|
||
workers int
|
||
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, workers int) (*Detector, error) {
|
||
if threshold <= 0 || threshold >= 1 {
|
||
threshold = 0.5
|
||
}
|
||
if workers <= 0 {
|
||
workers = 2 // 默认 2 个 worker
|
||
}
|
||
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,
|
||
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
|
||
}
|
||
|
||
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.Env = append(os.Environ(), fmt.Sprintf("KM_ANIME_NSFW_WORKERS=%d", d.workers))
|
||
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
|
||
}
|