327 lines
13 KiB
Python
327 lines
13 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""NudeNet 本地检测服务：由 Go 侧启动，常驻提供 /health 和 /detect。
|
||
|
||
只依赖 Python 标准库 + nudenet 包。
|
||
onnxruntime 的 InferenceSession 不是线程安全的，同一个实例并发推理会崩，
|
||
所以这里用一个固定大小的推理线程池：每个 worker 线程各持一份独立模型实例，
|
||
N 个请求可以真·并行推理。收到 /detect 就把图片喂给 NudeDetector，返回结果 JSON。
|
||
|
||
模型加载优先级（模型放在 <程序运行目录>/lib/nwfs/320n.onnx）：
|
||
1. 本地已有该文件 → 直接加载（离线可用，适合随程序分发）；
|
||
2. 本地没有 → 自动下载到 lib/nwfs/320n.onnx 再加载（仅第一次需要网络）；
|
||
3. 下载失败 → 回退用 nudenet 包自带的模型。
|
||
下载地址是 NudeNet 官方 release 的 320n.onnx（约 11.6MB）。
|
||
|
||
Windows 超长路径说明：OpenCV 的 imread 走老式 C API，路径超过 260 字符就
|
||
打不开（报 can't open/read file）。所以这里改用 Python 的 open() 读字节，
|
||
再把字节直接交给 NudeNet（v3 的 detect 支持 bytes）——从内存解码，彻底
|
||
绕开路径限制，同时兼容老版 v2（只收路径时落到临时短路径文件再试）。
|
||
"""
|
||
import concurrent.futures
|
||
import json
|
||
import os
|
||
import shutil
|
||
import sys
|
||
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:
|
||
cv2.setLogLevel(8) # LOG_LEVEL_ERROR
|
||
except Exception:
|
||
pass
|
||
|
||
# 只有这些"暴露部位"才算色情；其它如腋窝/脚踝暴露不判。
|
||
SENSITIVE = {
|
||
"FEMALE_GENITALIA_EXPOSED", # 女性生殖器暴露
|
||
"MALE_GENITALIA_EXPOSED", # 男性生殖器暴露
|
||
"ANUS_EXPOSED", # 肛门暴露
|
||
"FEMALE_BREAST_EXPOSED", # 女性乳房暴露
|
||
"BUTTOCKS_EXPOSED", # 臀部暴露
|
||
"BELLY_EXPOSED", # 腹部/腰部暴露（比基尼、露脐装等）
|
||
"FEMALE_GENITALIA_COVERED", # 女性生殖器被遮挡（透明布/薄纱等）
|
||
}
|
||
|
||
# 模型文件与官方下载地址
|
||
MODEL_FILENAME = "nwfs/320n.onnx"
|
||
MODEL_URL = "https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/320n.onnx"
|
||
|
||
# 模型目录由 Go 侧作为第二个参数传入（程序运行目录下的 lib）；没传则回退用包自带模型
|
||
MODEL_DIR = sys.argv[2] if len(sys.argv) > 2 else ""
|
||
|
||
_detector = None
|
||
|
||
# 推理并发数：onnxruntime 的 InferenceSession.run() 不是线程安全的，
|
||
# 同一个实例并发推理会崩，所以每个 worker 线程各持有一个独立实例（NudeDetector），
|
||
# N 个 worker 即可真·并行推理。默认 4，可用环境变量 KM_NSFW_WORKERS 调整。
|
||
INFER_WORKERS = int(os.environ.get("KM_NSFW_WORKERS", "4") or "4")
|
||
if INFER_WORKERS < 1:
|
||
INFER_WORKERS = 1
|
||
|
||
_executor = None # 推理线程池（固定 N 个 worker，模型各持一份）
|
||
_thread_local = threading.local() # 每个 worker 线程自己的检测器
|
||
_MODEL_PATH = None # 启动时解析好的模型路径（None 表示用包自带模型）
|
||
|
||
|
||
def _looks_like_model(path):
|
||
"""粗略校验文件是不是真 ONNX 模型：开头不是 HTML、大小不小于 1MB（320n 模型约 11.6MB）。
|
||
|
||
防止 GitHub 被墙时把"登录网页/占位页"当成模型存下来，下次加载直接报错。
|
||
"""
|
||
try:
|
||
with open(path, "rb") as f:
|
||
head = f.read(4)
|
||
if not head or head.startswith(b"<"):
|
||
return False
|
||
return os.path.getsize(path) > 1_000_000
|
||
except OSError:
|
||
return False
|
||
|
||
|
||
def _download_model(local):
|
||
"""带超时地把模型下载到 local。urlretrieve 没有超时，网络挂起会卡死整个服务。"""
|
||
req = urllib.request.Request(MODEL_URL)
|
||
with urllib.request.urlopen(req, timeout=30) as resp, open(local, "wb") as f:
|
||
shutil.copyfileobj(resp, f)
|
||
|
||
|
||
def _ensure_model(model_dir):
|
||
"""返回要加载的模型路径；返回 None 表示用 nudenet 包自带的模型。"""
|
||
if not model_dir:
|
||
return None
|
||
local = os.path.join(model_dir, MODEL_FILENAME)
|
||
if os.path.isfile(local) and _looks_like_model(local):
|
||
print("=== 使用本地模型: %s" % local, flush=True)
|
||
return local
|
||
# 本地没有（或文件无效）→ 下载到 lib 目录（仅第一次需要网络，带 30 秒超时）
|
||
try:
|
||
os.makedirs(model_dir, exist_ok=True)
|
||
print("=== 本地未找到有效模型, 正在下载: %s" % local, flush=True)
|
||
_download_model(local)
|
||
if _looks_like_model(local):
|
||
print("=== 模型下载完成: %s" % local, flush=True)
|
||
return local
|
||
# 下载回来的是网页/占位文件，不是真模型：删掉免得下次继续踩坑
|
||
print("=== 下载的文件无效(可能被墙返回了网页), 删除并回退用包自带模型", flush=True)
|
||
try:
|
||
os.remove(local)
|
||
except OSError:
|
||
pass
|
||
except Exception as exc: # noqa: BLE001 - 下载失败就回退，别让服务起不来
|
||
print("=== 模型下载失败, 回退用包自带模型: %s" % exc, flush=True)
|
||
return None
|
||
|
||
|
||
def get_detector():
|
||
"""确保模型已就绪（下载/加载），返回一个检测器实例。
|
||
|
||
启动时调用一次：把模型下载/解析好，/health 就绪即代表模型可用。
|
||
真正推理用的独立实例由 _new_detector() 按线程各建一份，共享这份模型路径。
|
||
"""
|
||
global _detector, _MODEL_PATH
|
||
if _detector is None:
|
||
_MODEL_PATH = _ensure_model(MODEL_DIR)
|
||
_detector = _new_detector()
|
||
return _detector
|
||
|
||
|
||
def _new_detector():
|
||
"""创建一份独立检测器实例（每个推理 worker 各持一份，可并行推理）。"""
|
||
mp = _MODEL_PATH
|
||
try:
|
||
return NudeDetector(model_path=mp) if mp else NudeDetector()
|
||
except Exception as exc: # noqa: BLE001 - 指定路径加载失败就回退包自带模型，别崩
|
||
print("=== 本地模型加载失败(%s), 回退用包自带模型" % exc, flush=True)
|
||
if mp:
|
||
try:
|
||
os.remove(mp)
|
||
except OSError:
|
||
pass
|
||
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)
|
||
|
||
# 没有检测到任何类别，直接返回空
|
||
if not boxes:
|
||
return []
|
||
|
||
# 转成 numpy array，OpenCV 的 NMSBoxes 需要
|
||
boxes_arr = np.array(boxes, dtype=np.float32)
|
||
scores_arr = np.array(scores, dtype=np.float32)
|
||
|
||
# NMS 参数：score_threshold=0.01 保留更多低分数结果，iou_threshold=0.45
|
||
indices = cv2.dnn.NMSBoxes(boxes_arr.tolist(), scores_arr.tolist(), 0.01, 0.45)
|
||
|
||
detections = []
|
||
if indices is not None and len(indices) > 0:
|
||
for i in np.array(indices).flatten():
|
||
idx = int(i)
|
||
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":
|
||
return path
|
||
if path.startswith("\\\\?\\"):
|
||
return path
|
||
if path.startswith("\\\\"):
|
||
# UNC 网络路径：\\server\share → \\?\UNC\server\share
|
||
return "\\\\?\\UNC\\" + path.lstrip("\\")
|
||
if os.path.isabs(path):
|
||
return "\\\\?\\" + os.path.normpath(path)
|
||
return path
|
||
|
||
|
||
def _read_bytes(path):
|
||
"""用 Python 读文件字节（不经过 OpenCV），读不了就抛异常由上层兜底。"""
|
||
lp = _longpath(path)
|
||
if not os.path.exists(lp):
|
||
raise OSError("文件不存在: %s" % path)
|
||
with open(lp, "rb") as f:
|
||
data = f.read()
|
||
if not data:
|
||
raise OSError("空文件: %s" % path)
|
||
return data
|
||
|
||
|
||
def _infer(path):
|
||
"""整个推理任务（读文件 + 检测），提交到固定线程池执行。
|
||
|
||
使用 _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 _raw_detect(det, data)
|
||
|
||
|
||
class Handler(BaseHTTPRequestHandler):
|
||
def _send(self, obj, status=200):
|
||
data = json.dumps(obj, ensure_ascii=False).encode("utf-8")
|
||
self.send_response(status)
|
||
self.send_header("Content-Type", "application/json; charset=utf-8")
|
||
self.send_header("Content-Length", str(len(data)))
|
||
self.end_headers()
|
||
self.wfile.write(data)
|
||
|
||
def do_GET(self):
|
||
if self.path == "/health":
|
||
self._send({"ok": True})
|
||
else:
|
||
self._send({"ok": False, "error": "not found"}, 404)
|
||
|
||
def do_POST(self):
|
||
if self.path != "/detect":
|
||
return self._send({"ok": False, "error": "not found"}, 404)
|
||
length = int(self.headers.get("Content-Length", 0))
|
||
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 并行处理
|
||
detections = _executor.submit(_infer, path).result()
|
||
except Exception as exc: # noqa: BLE001 - 服务端任何异常都要回给 Go，别让进程崩
|
||
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})
|
||
|
||
def log_message(self, *args): # 不刷屏
|
||
pass
|
||
|
||
|
||
# 多线程接收连接 + 大监听队列：Go 侧会并发发请求（worker 数可到 8），
|
||
# 单线程 HTTPServer 默认 backlog 只有 5，并发一高连接就被内核拒绝（报 10061 连接拒绝）。
|
||
# ThreadingHTTPServer 让连接即时被 accept；推理再交给固定线程池并行执行
|
||
# （每个 worker 一份独立模型实例，onnxruntime 同一实例并发推理会崩）。
|
||
class Server(ThreadingHTTPServer):
|
||
request_queue_size = 128
|
||
|
||
|
||
if __name__ == "__main__":
|
||
port = int(sys.argv[1]) if len(sys.argv) > 1 else 0
|
||
# 启动前先把模型准备好（本地加载 / 下载 / 回退），/health 就绪即代表模型可用，
|
||
# 这样 Go 侧 waitReady 等到的就是"模型已就绪"的服务。
|
||
get_detector()
|
||
# 固定大小推理线程池：每个 worker 线程持一份独立模型实例，可并行推理
|
||
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS)
|
||
Server(("127.0.0.1", port), Handler).serve_forever()
|