upd

internal/nsfw/detector.go

@@ -39,6 +39,13 @@ 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
@@ -210,28 +217,29 @@ func (d *Detector) waitReady() error {
 // Detect 检测一张图片是否色情。返回：
 //   - score：命中的敏感类里最高分数（没有命中为 0）
 //   - nsfw：是否判定为色情（有敏感类且分数 >= 阈值）
+//   - detections：所有检测到的目标（类别+分数+位置）
 //   - err：检测失败（服务挂了/图读不出等），由调用方兜底按未命中处理
 //
 // 服务失联（连接被拒，多半是 python 进程在处理某张图时崩了）时，会串行化地
 // 重启服务并重试一次：多个并发 worker 里只有一个负责重启，其余排队，等重启完
 // 也重试自己的图——保证服务短暂挂掉不会漏掉一整批图片。
-func (d *Detector) Detect(path string) (score float64, nsfw bool, err error) {
-	score, nsfw, err = d.detectOnce(path)
+func (d *Detector) Detect(path string) (score float64, nsfw bool, detections []NsfwDetection, err error) {
+	score, nsfw, detections, err = d.detectOnce(path)
 	if err == nil {
-		return score, nsfw, nil
+		return score, nsfw, detections, nil
 	}
 	// 只有"发请求都失败"（连接被拒/重置，说明服务失联）才走重启；
 	// 服务正常返回的明确错误（如某张图读不了）不重启。
 	var ce *connError
 	if !errors.As(err, &ce) {
-		return score, nsfw, err
+		return score, nsfw, detections, err
 	}
 
 	d.restartMu.Lock()
 	defer d.restartMu.Unlock()
 
 	if d.isStopped() {
-		return 0, false, err
+		return 0, false, nil, err
 	}
 	// 可能排在前面的 goroutine 已经把服务重启好了：先探活，活着就直接重试，别反复重启
 	if d.healthy() {
@@ -239,7 +247,7 @@ func (d *Detector) Detect(path string) (score float64, nsfw bool, err error) {
 	}
 	fmt.Printf("=== 色情检测服务失联(%v), 正在重启…\n", ce.err)
 	if rerr := d.launch(); rerr != nil {
-		return 0, false, fmt.Errorf("检测服务重启失败: %v", rerr)
+		return 0, false, nil, fmt.Errorf("检测服务重启失败: %v", rerr)
 	}
 	fmt.Println("=== 色情检测服务已重启，继续识别")
 	return d.detectOnce(path)
@@ -269,7 +277,7 @@ func (d *Detector) isStopped() bool {
 }
 
 // detectOnce 发一次检测请求（不带重启逻辑）。
-func (d *Detector) detectOnce(path string) (score float64, nsfw bool, err error) {
+func (d *Detector) detectOnce(path string) (score float64, nsfw bool, detections []NsfwDetection, err error) {
 	body, _ := json.Marshal(map[string]interface{}{
 		"path":      path,
 		"threshold": d.threshold,
@@ -277,12 +285,12 @@ func (d *Detector) detectOnce(path string) (score float64, nsfw bool, err error)
 	resp, err := d.client.Post(d.url+"/detect", "application/json", bytes.NewReader(body))
 	if err != nil {
 		// 打上"连接失败"标记，供 Detect 判断是否要重启
-		return 0, false, &connError{err}
+		return 0, false, nil, &connError{err}
 	}
 	defer resp.Body.Close()
 	data, err := io.ReadAll(resp.Body)
 	if err != nil {
-		return 0, false, err
+		return 0, false, nil, err
 	}
 	var out struct {
 		Ok         bool                     `json:"ok"`
@@ -292,24 +300,36 @@ func (d *Detector) detectOnce(path string) (score float64, nsfw bool, err error)
 		Detections []map[string]interface{} `json:"detections"`
 	}
 	if err := json.Unmarshal(data, &out); err != nil {
-		return 0, false, fmt.Errorf("解析检测结果失败: %v", err)
+		return 0, false, nil, fmt.Errorf("解析检测结果失败: %v", err)
 	}
 	if !out.Ok {
-		return 0, false, fmt.Errorf("检测服务返回错误: %s", out.Error)
+		return 0, false, 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(out.Detections) > 0 {
+	if len(detections) > 0 {
 		var classes []string
-		for _, det := range out.Detections {
-			cls, _ := det["class"].(string)
-			sc, _ := det["score"].(float64)
-			classes = append(classes, fmt.Sprintf("%s:%.3f", cls, sc))
+		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, out.Nsfw, nil
+	return out.Score, out.Nsfw, detections, nil
 }
 
 // connError 标记"请求根本没发出去/连不上本地服务"的错误（服务失联）。