diff --git a/README.md b/README.md index 7fc05b6..157045f 100644 --- a/README.md +++ b/README.md @@ -187,7 +187,8 @@ lightScore = 亮度 > 175 的像素占比 基于本机 **Python + NudeNet** 的离线鉴黄,判为色情的图片移入 `图片_敏感_*` 分类目录(明暗子目录照常保留,如 `图片_敏感_横屏_1080P/bright_20`),便于事后人工复核。 - **开启**:`application.properties` 里设 `nsfw.enabled=true`(默认关闭) -- **依赖**:本机安装 Python(>=3.6)并 `pip install nudenet`(包自带 320n 模型,无需手动下载) +- **依赖**:本机安装 Python 并 `pip install nudenet` +- **模型(推荐分发方式)**:把模型文件放到程序运行目录的 `lib/320n.onnx`(官方地址 `https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/320n.onnx`,约 11.6MB,需自带或下载后随程序分发)——放好后完全离线可用。`lib` 里没有时程序会自动下载一次到该位置;下载失败再回退用 nudenet 包自带的模型 - **可调**:`nsfw.python`(python 路径,留空自动在 PATH 找)、`nsfw.threshold`(判定阈值,默认 0.5) - **敏感类**:`FEMALE_GENITALIA_EXPOSED` / `MALE_GENITALIA_EXPOSED` / `ANUS_EXPOSED` / `FEMALE_BREAST_EXPOSED`,任一分数达标即判定 - **降级**:未装 python / nudenet 时打印警告并跳过识别,图片按普通流程整理,不影响其它功能 diff --git a/application.properties b/application.properties index ad86ea9..8ebb595 100644 --- a/application.properties +++ b/application.properties @@ -12,8 +12,10 @@ ffmpeg.path= organize.rotateImages=false # 探测/认类型/移动阶段的并发 worker 数(默认 4;0 表示用满全部 CPU 核数) organize.workers=4 -# 色情图片识别(需本机安装 python 与 nudenet:pip install nudenet),需要科技下载模型 +# 色情图片识别(需本机安装 python 与 nudenet:pip install nudenet) # true=开启,判为色情的图片移入「图片_敏感_*」目录(明暗子目录照常);false=关闭(默认) +# 模型:优先加载程序运行目录下 lib/320n.onnx(推荐随程序分发,离线可用); +# 本地没有会自动下载到该位置,下载失败再回退用 nudenet 自带模型 nsfw.enabled=true # python 可执行文件路径;留空自动在 PATH 里找 python/python3/py nsfw.python=C:\\Users\\Administrator\\AppData\\Local\\Programs\\Python\\Python312\\python.exe diff --git a/internal/nsfw/detector.go b/internal/nsfw/detector.go index 40d6df7..122b2f1 100644 --- a/internal/nsfw/detector.go +++ b/internal/nsfw/detector.go @@ -39,8 +39,8 @@ var serverScript []byte // Go 侧只负责把路径发给服务、读取判定结果,不重复维护这份清单。 // detectorStartTimeout 等待 python 服务就绪(模型加载)的超时时间。 -// 首次加载模型可能要十几秒,给足余量。 -const detectorStartTimeout = 60 * time.Second +// 首次启动若本地没有模型还要联网下载约 11.6MB,给足余量。 +const detectorStartTimeout = 180 * time.Second // requestTimeout 单张图片检测的超时时间。检测很快(毫秒级), // 万一卡住(模型没加载好/图太大)不至于拖死整个扫描。 @@ -94,9 +94,16 @@ func Start(pythonPath string, threshold float64) (*Detector, error) { return nil, fmt.Errorf("找空闲端口失败: %v", err) } + // 模型目录:程序运行目录下的 lib(和 application.properties / media-info.db 同一根目录)。 + // 服务脚本优先加载 lib/320n.onnx;本地没有会自己下载到那里,下载失败再回退用 pip 自带模型。 + modelDir := "" + if wd, err := os.Getwd(); err == nil { + modelDir = filepath.Join(wd, "lib") + } + // 拉起服务。标准输出/错误接到程序自身的输出,方便看 python 报错。 ctx, cancel := context.WithCancel(context.Background()) - cmd := exec.CommandContext(ctx, py, scriptPath, fmt.Sprintf("%d", port)) + cmd := exec.CommandContext(ctx, py, scriptPath, fmt.Sprintf("%d", port), modelDir) cmd.Stdout = os.Stdout cmd.Stderr = os.Stderr if err := cmd.Start(); err != nil { diff --git a/internal/nsfw/server.py b/internal/nsfw/server.py index 0a2ad0e..0df3a49 100644 --- a/internal/nsfw/server.py +++ b/internal/nsfw/server.py @@ -4,6 +4,12 @@ 只依赖 Python 标准库 + nudenet 包,模型在进程里只加载一次。 收到 /detect 就把图片喂给 NudeDetector,返回结果 JSON。 +模型加载优先级(模型放在 <程序运行目录>/lib/320n.onnx): + 1. 本地已有该文件 → 直接加载(离线可用,适合随程序分发); + 2. 本地没有 → 自动下载到 lib/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)——从内存解码,彻底 @@ -13,6 +19,7 @@ import json import os import sys import tempfile +import urllib.request from http.server import BaseHTTPRequestHandler, HTTPServer from nudenet import NudeDetector @@ -33,13 +40,72 @@ SENSITIVE = { "FEMALE_BREAST_EXPOSED", } +# 模型文件与官方下载地址 +MODEL_FILENAME = "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 +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 _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 目录(仅第一次需要网络) + try: + os.makedirs(model_dir, exist_ok=True) + print("=== 本地未找到有效模型, 正在下载: %s" % local, flush=True) + urllib.request.urlretrieve(MODEL_URL, 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(): global _detector if _detector is None: - _detector = NudeDetector() + mp = _ensure_model(MODEL_DIR) + try: + _detector = NudeDetector(model_path=mp) + except Exception as exc: # noqa: BLE001 - 指定路径加载失败就回退包自带模型,别崩 + print("=== 本地模型加载失败(%s), 回退用包自带模型" % exc, flush=True) + if mp: + try: + os.remove(mp) + except OSError: + pass + _detector = NudeDetector() return _detector @@ -136,4 +202,7 @@ class Handler(BaseHTTPRequestHandler): if __name__ == "__main__": port = int(sys.argv[1]) if len(sys.argv) > 1 else 0 + # 启动前先把模型准备好(本地加载 / 下载 / 回退),/health 就绪即代表模型可用, + # 这样 Go 侧 waitReady 等到的就是"模型已就绪"的服务。 + get_detector() HTTPServer(("127.0.0.1", port), Handler).serve_forever() \ No newline at end of file diff --git a/lib/320n.onnx b/lib/320n.onnx new file mode 100644 index 0000000..68eedc3 Binary files /dev/null and b/lib/320n.onnx differ