feat: 降低 NSFW 检测阈值，显示更多检测结果
This commit is contained in:
@@ -306,6 +306,8 @@ func (d *Detector) detectOnce(path string) (score float64, nsfw bool, err error)
|
||||
classes = append(classes, fmt.Sprintf("%s:%.3f", cls, sc))
|
||||
}
|
||||
fmt.Printf("=== NSFW 检测详情 %s: [%s]\n", path, strings.Join(classes, ", "))
|
||||
} else {
|
||||
fmt.Printf("=== NSFW 检测详情 %s: [无检测结果]\n", path)
|
||||
}
|
||||
return out.Score, out.Nsfw, nil
|
||||
}
|
||||
|
||||
+73
-27
@@ -22,17 +22,17 @@ import json
|
||||
import os
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import urllib.request
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from nudenet import NudeDetector
|
||||
from nudenet.nudenet import _read_image
|
||||
|
||||
# 只报 OpenCV 致命错误，别让 imread 的 WARN 刷屏（也避免误导排查）
|
||||
try:
|
||||
import cv2
|
||||
|
||||
cv2.setLogLevel(8) # LOG_LEVEL_ERROR
|
||||
except Exception:
|
||||
pass
|
||||
@@ -145,6 +145,74 @@ def _new_detector():
|
||||
return NudeDetector()
|
||||
|
||||
|
||||
def _raw_detect(det, data):
|
||||
"""直接调用 ONNX 推理，用 0.05 低阈值后处理（NudeNet 内置的 detect 用 0.2 太严格）。
|
||||
|
||||
绕过 NudeDetector.detect() 的硬编码 0.2 阈值，让低置信度的检测结果也能返回，
|
||||
最终判定交给 Go 侧的用户配置阈值。
|
||||
"""
|
||||
mat = cv2.imdecode(np.frombuffer(data, np.uint8), -1)
|
||||
if mat is None:
|
||||
raise OSError("图片解码失败")
|
||||
|
||||
mat_c3 = cv2.cvtColor(mat, cv2.COLOR_RGBA2BGR)
|
||||
ow, oh = mat_c3.shape[1], mat_c3.shape[0]
|
||||
max_size = max(ow, oh)
|
||||
x_pad = max_size - ow
|
||||
y_pad = max_size - oh
|
||||
x_ratio = max_size / ow
|
||||
y_ratio = max_size / oh
|
||||
|
||||
mat_pad = cv2.copyMakeBorder(mat_c3, 0, y_pad, 0, x_pad, cv2.BORDER_CONSTANT)
|
||||
blob = cv2.dnn.blobFromImage(mat_pad, 1 / 255.0, (det.input_width, det.input_height), (0, 0, 0), swapRB=True, crop=False)
|
||||
|
||||
outputs = det.onnx_session.run(None, {det.input_name: blob})
|
||||
outs = np.transpose(np.squeeze(outputs[0]))
|
||||
rows = outs.shape[0]
|
||||
|
||||
boxes, scores, class_ids = [], [], []
|
||||
for i in range(rows):
|
||||
cls_scores = outs[i][4:]
|
||||
max_score = np.amax(cls_scores)
|
||||
if max_score >= 0.05: # 低阈值：NudeNet 内置是 0.2，这里降到 0.05
|
||||
class_id = int(np.argmax(cls_scores))
|
||||
x, y, w, h = outs[i][0:4]
|
||||
x, y = x - w / 2, y - h / 2
|
||||
w_ = w * (ow + x_pad) / det.input_width
|
||||
h_ = h * (oh + y_pad) / det.input_height
|
||||
x = x * (ow + x_pad) / det.input_width
|
||||
y = y * (oh + y_pad) / det.input_height
|
||||
x, y = max(0, min(x, ow)), max(0, min(y, oh))
|
||||
w_ = min(w_, ow - x)
|
||||
h_ = min(h_, oh - y)
|
||||
boxes.append([x, y, w_, h_])
|
||||
scores.append(float(max_score))
|
||||
class_ids.append(class_id)
|
||||
|
||||
indices = cv2.dnn.NMSBoxes(boxes, scores, 0.1, 0.45)
|
||||
|
||||
detections = []
|
||||
for i in indices:
|
||||
idx = i if isinstance(i, int) else i[0]
|
||||
detections.append({
|
||||
"class": __labels[class_ids[idx]],
|
||||
"score": scores[idx],
|
||||
"box": [int(boxes[idx][0]), int(boxes[idx][1]), int(boxes[idx][2]), int(boxes[idx][3])],
|
||||
})
|
||||
return detections
|
||||
|
||||
|
||||
# NudeNet 的标签表（和 nudenet.nudenet.__labels 一致）
|
||||
__labels = [
|
||||
"FEMALE_GENITALIA_COVERED", "FACE_FEMALE", "BUTTOCKS_EXPOSED",
|
||||
"FEMALE_BREAST_EXPOSED", "FEMALE_GENITALIA_EXPOSED", "MALE_BREAST_EXPOSED",
|
||||
"ANUS_EXPOSED", "FEET_EXPOSED", "BELLY_COVERED", "FEET_COVERED",
|
||||
"ARMPITS_COVERED", "ARMPITS_EXPOSED", "FACE_MALE", "BELLY_EXPOSED",
|
||||
"MALE_GENITALIA_EXPOSED", "ANUS_COVERED", "FEMALE_BREAST_COVERED",
|
||||
"BUTTOCKS_COVERED",
|
||||
]
|
||||
|
||||
|
||||
def _longpath(path):
|
||||
"""Windows 下给绝对路径加 \\\\?\\ 前缀，让 Python 原生 open 也能处理超长路径。"""
|
||||
if os.name != "nt":
|
||||
@@ -171,39 +239,17 @@ def _read_bytes(path):
|
||||
return data
|
||||
|
||||
|
||||
def _detect_bytes(det, data, path):
|
||||
"""把图片字节喂给 NudeNet。v3 直接吃 bytes；老版 v2 只收路径时落到临时短路径重试。
|
||||
|
||||
用线程自己的检测器实例推理，各实例互不干扰，可以真·并行。
|
||||
"""
|
||||
try:
|
||||
return det.detect(data)
|
||||
except (TypeError, AttributeError):
|
||||
tmp = None
|
||||
try:
|
||||
fd, tmp = tempfile.mkstemp(suffix=os.path.splitext(path)[1] or ".jpg")
|
||||
with os.fdopen(fd, "wb") as f:
|
||||
f.write(data)
|
||||
return det.detect(tmp)
|
||||
finally:
|
||||
if tmp:
|
||||
try:
|
||||
os.unlink(tmp)
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _infer(path):
|
||||
"""整个推理任务（读文件 + 检测），提交到固定线程池执行。
|
||||
|
||||
读文件在池外并行进行；推理用线程各自的独立检测器，N 个 worker 真·并行。
|
||||
使用 _raw_detect 绕过 NudeNet 内置的 0.2 阈值，用 0.05 低阈值返回更多结果。
|
||||
"""
|
||||
data = _read_bytes(path)
|
||||
det = getattr(_thread_local, "detector", None)
|
||||
if det is None:
|
||||
det = _new_detector()
|
||||
_thread_local.detector = det
|
||||
return _detect_bytes(det, data, path)
|
||||
return _raw_detect(det, data)
|
||||
|
||||
|
||||
class Handler(BaseHTTPRequestHandler):
|
||||
|
||||
Reference in New Issue
Block a user