77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""NudeNet 本地检测服务：由 Go 侧启动，常驻提供 /health 和 /detect。
|
|
|
|
只依赖 Python 标准库 + nudenet 包，模型在进程里只加载一次。
|
|
收到 /detect 就把图片路径丢给 NudeDetector，返回结果 JSON。
|
|
"""
|
|
import json
|
|
import sys
|
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
|
|
|
from nudenet import NudeDetector
|
|
|
|
# 只有这些"暴露部位"才算色情；其它如肚子/腋窝/脚踝暴露不判。
|
|
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
|
|
|
|
|
|
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))
|
|
detections = get_detector().detect(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()
|