odboy/kenaito-media-picker
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 目录层级变更

Browse Source
This commit is contained in:
odboy
2025-09-24 21:19:55 +08:00
parent 7792f9281c
commit 72c9d67d9b
12 changed files with 156 additions and 317 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/nsfw/detector.go
+17 -31
View File
@@ -58,9 +58,8 @@ const requestTimeout = 30 * time.Second
// 多个探测协程会并发调 Detect，底层 HTTP 客户端线程安全；
// 服务端用固定大小线程池并行推理（每个 worker 一份独立模型实例），请求不必排队。
type Detector struct {
url string // 本地服务地址，如 http://127.0.0.1:54321
threshold float64 // 判定阈值 0~1
workers int // Python 推理 worker 数量
url string // 本地服务地址，如 http://127.0.0.1:54321
workers int // Python 推理 worker 数量
client *http.Client
cmd *exec.Cmd // python 进程句柄（Stop 时杀掉）
scriptDir string // 临时脚本所在目录（Stop 时清理）
@@ -75,14 +74,10 @@ type Detector struct {
// Start 拉起本地 python 检测服务并等它就绪。
// - pythonPath：python 可执行文件路径；为空时在 PATH 里找 python/python3。
// - threshold：判定阈值 0~1。
// - workers：Python 推理 worker 数量；<=0 表示使用默认值（4）。
//
// 找不到 python、装不了 nudenet、或服务起不来时返回错误（调用方跳过识别即可）。
func Start(pythonPath string, threshold float64, workers int) (*Detector, error) {
if threshold <= 0 || threshold >= 1 {
threshold = 0.5 // 阈值传得不合理就用默认
}
func Start(pythonPath string, workers int) (*Detector, error) {
if workers <= 0 {
workers = 4 // 默认 4 个 worker
}
@@ -111,7 +106,6 @@ func Start(pythonPath string, threshold float64, workers int) (*Detector, error)
}
d := &Detector{
threshold: threshold,
workers: workers,
client: &http.Client{Timeout: requestTimeout},
python: py,
@@ -222,32 +216,31 @@ func (d *Detector) waitReady() error {
return fmt.Errorf("检测服务在 %v 内未就绪", detectorStartTimeout)
}
// Detect 检测一张图片是否色情。返回：
// Detect 检测一张图片的色情评分。返回：
// - score：命中的敏感类里最高分数（没有命中为 0）
// - nsfw：是否判定为色情（有敏感类且分数 >= 阈值）
// - detections：所有检测到的目标（类别+分数+位置）
// - err：检测失败（服务挂了/图读不出等），由调用方兜底按未命中处理
//
// 服务失联（连接被拒，多半是 python 进程在处理某张图时崩了）时，会串行化地
// 重启服务并重试一次：多个并发 worker 里只有一个负责重启，其余排队，等重启完
// 也重试自己的图——保证服务短暂挂掉不会漏掉一整批图片。
func (d *Detector) Detect(path string) (score float64, nsfw bool, detections []NsfwDetection, err error) {
score, nsfw, detections, err = d.detectOnce(path)
func (d *Detector) Detect(path string) (score float64, detections []NsfwDetection, err error) {
score, detections, err = d.detectOnce(path)
if err == nil {
return score, nsfw, detections, nil
return score, detections, nil
}
// 只有"发请求都失败"（连接被拒/重置，说明服务失联）才走重启；
// 服务正常返回的明确错误（如某张图读不了）不重启。
var ce *connError
if !errors.As(err, &ce) {
return score, nsfw, detections, err
return score, detections, err
}
d.restartMu.Lock()
defer d.restartMu.Unlock()
if d.isStopped() {
return 0, false, nil, err
return 0, nil, err
}
// 可能排在前面的 goroutine 已经把服务重启好了：先探活，活着就直接重试，别反复重启
if d.healthy() {
@@ -255,7 +248,7 @@ func (d *Detector) Detect(path string) (score float64, nsfw bool, detections []N
}
fmt.Printf("=== 色情检测服务失联(%v), 正在重启…\n", ce.err)
if rerr := d.launch(); rerr != nil {
return 0, false, nil, fmt.Errorf("检测服务重启失败: %v", rerr)
return 0, nil, fmt.Errorf("检测服务重启失败: %v", rerr)
}
fmt.Println("=== 色情检测服务已重启，继续识别")
return d.detectOnce(path)
@@ -285,33 +278,31 @@ func (d *Detector) isStopped() bool {
}
// detectOnce 发一次检测请求（不带重启逻辑）。
func (d *Detector) detectOnce(path string) (score float64, nsfw bool, detections []NsfwDetection, err error) {
func (d *Detector) detectOnce(path string) (score float64, detections []NsfwDetection, err error) {
body, _ := json.Marshal(map[string]interface{}{
"path": path,
"threshold": d.threshold,
"path": path,
})
resp, err := d.client.Post(d.url+"/detect", "application/json", bytes.NewReader(body))
if err != nil {
// 打上"连接失败"标记，供 Detect 判断是否要重启
return 0, false, nil, &connError{err}
return 0, nil, &connError{err}
}
defer resp.Body.Close()
data, err := io.ReadAll(resp.Body)
if err != nil {
return 0, false, nil, err
return 0, nil, err
}
var out struct {
Ok bool `json:"ok"`
Error string `json:"error"`
Nsfw bool `json:"nsfw"`
Score float64 `json:"score"`
Detections []map[string]interface{} `json:"detections"`
}
if err := json.Unmarshal(data, &out); err != nil {
return 0, false, nil, fmt.Errorf("解析检测结果失败: %v", err)
return 0, nil, fmt.Errorf("解析检测结果失败: %v", err)
}
if !out.Ok {
return 0, false, nil, fmt.Errorf("检测服务返回错误: %s", out.Error)
return 0, nil, fmt.Errorf("检测服务返回错误: %s", out.Error)
}
// 解析检测详情
detections = make([]NsfwDetection, 0, len(out.Detections))
@@ -337,7 +328,7 @@ func (d *Detector) detectOnce(path string) (score float64, nsfw bool, detections
} else {
fmt.Printf("=== NSFW 检测详情 %s: [无检测结果]\n", path)
}
return out.Score, out.Nsfw, detections, nil
return out.Score, detections, nil
}
// connError 标记"请求根本没发出去/连不上本地服务"的错误（服务失联）。
@@ -360,11 +351,6 @@ func (d *Detector) Stop() {
}
}
// 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 {
internal/nsfw/server.py
+1 -5
View File
@@ -284,7 +284,6 @@ class Handler(BaseHTTPRequestHandler):
try:
body = json.loads(self.rfile.read(length) or b"{}")
path = body.get("path", "")
threshold = float(body.get("threshold", 0.5))
if not path:
return self._send({"ok": False, "error": "缺少图片路径"})
# 丢给推理线程池：读文件并行、推理由 N 个 worker 并行处理
@@ -293,16 +292,13 @@ class Handler(BaseHTTPRequestHandler):
return self._send({"ok": False, "error": str(exc)})
max_score = 0.0
nsfw = False
for det in detections or []:
cls = det.get("class", "")
sc = float(det.get("score", 0) or 0)
if cls in SENSITIVE:
if sc > max_score:
max_score = sc
if sc >= threshold:
nsfw = True
self._send({"ok": True, "nsfw": nsfw, "score": max_score, "detections": detections})
self._send({"ok": True, "score": max_score, "detections": detections})
def log_message(self, *args): # 不刷屏
pass