feat: 二次元识别功能 + 模型目录整理(nwfs/wd-tagger 子目录),修复 wd-tagger 输入维度(NHWC)

This commit is contained in:
2025-06-25 21:10:21 +08:00
parent 4efa48f550
commit 3af2bf70b4
18 changed files with 1386 additions and 21 deletions
+4 -4
View File
@@ -6,9 +6,9 @@ onnxruntime 的 InferenceSession 不是线程安全的,同一个实例并发
所以这里用一个固定大小的推理线程池:每个 worker 线程各持一份独立模型实例,
N 个请求可以真·并行推理。收到 /detect 就把图片喂给 NudeDetector,返回结果 JSON。
模型加载优先级(模型放在 <程序运行目录>/lib/320n.onnx):
模型加载优先级(模型放在 <程序运行目录>/lib/nwfs/320n.onnx):
1. 本地已有该文件 → 直接加载(离线可用,适合随程序分发);
2. 本地没有 → 自动下载到 lib/320n.onnx 再加载(仅第一次需要网络);
2. 本地没有 → 自动下载到 lib/nwfs/320n.onnx 再加载(仅第一次需要网络);
3. 下载失败 → 回退用 nudenet 包自带的模型。
下载地址是 NudeNet 官方 release 的 320n.onnx(约 11.6MB)。
@@ -46,7 +46,7 @@ SENSITIVE = {
}
# 模型文件与官方下载地址
MODEL_FILENAME = "320n.onnx"
MODEL_FILENAME = "nwfs/320n.onnx"
MODEL_URL = "https://github.com/notAI-tech/NudeNet/releases/download/v3.4-weights/320n.onnx"
# 模型目录由 Go 侧作为第二个参数传入(程序运行目录下的 lib);没传则回退用包自带模型
@@ -265,4 +265,4 @@ if __name__ == "__main__":
get_detector()
# 固定大小推理线程池:每个 worker 线程持一份独立模型实例,可并行推理
_executor = concurrent.futures.ThreadPoolExecutor(max_workers=INFER_WORKERS)
Server(("127.0.0.1", port), Handler).serve_forever()
Server(("127.0.0.1", port), Handler).serve_forever()