139 lines
4.6 KiB
Python
139 lines
4.6 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""NudeNet 本地检测服务：由 Go 侧启动，常驻提供 /health 和 /detect。
|
||
|
||
只依赖 Python 标准库 + nudenet 包，模型在进程里只加载一次。
|
||
收到 /detect 就把图片喂给 NudeDetector，返回结果 JSON。
|
||
|
||
Windows 超长路径说明：OpenCV 的 imread 走老式 C API，路径超过 260 字符就
|
||
打不开（报 can't open/read file）。所以这里改用 Python 的 open() 读字节，
|
||
再把字节直接交给 NudeNet（v3 的 detect 支持 bytes）——从内存解码，彻底
|
||
绕开路径限制，同时兼容老版 v2（只收路径时落到临时短路径文件再试）。
|
||
"""
|
||
import json
|
||
import os
|
||
import sys
|
||
import tempfile
|
||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||
|
||
from nudenet import NudeDetector
|
||
|
||
# 只报 OpenCV 致命错误，别让 imread 的 WARN 刷屏（也避免误导排查）
|
||
try:
|
||
import cv2
|
||
|
||
cv2.setLogLevel(8) # LOG_LEVEL_ERROR
|
||
except Exception:
|
||
pass
|
||
|
||
# 只有这些"暴露部位"才算色情；其它如肚子/腋窝/脚踝暴露不判。
|
||
SENSITIVE = {
|
||
"FEMALE_GENITALIA_EXPOSED",
|
||
"MALE_GENITALIA_EXPOSED",
|
||
"ANUS_EXPOSED",
|
||
"FEMALE_BREAST_EXPOSED",
|
||
}
|
||
|
||
_detector = None
|
||
|
||
|
||
def get_detector():
|
||
global _detector
|
||
if _detector is None:
|
||
_detector = NudeDetector()
|
||
return _detector
|
||
|
||
|
||
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 _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
|
||
|
||
|
||
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": "缺少图片路径"})
|
||
data = _read_bytes(path)
|
||
detections = _detect_bytes(get_detector(), data, path)
|
||
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
|
||
|
||
|
||
if __name__ == "__main__":
|
||
port = int(sys.argv[1]) if len(sys.argv) > 1 else 0
|
||
HTTPServer(("127.0.0.1", port), Handler).serve_forever() |