feat: 二次元识别功能 + 模型目录整理(nwfs/wd-tagger 子目录),修复 wd-tagger 输入维度(NHWC)
This commit is contained in:
@@ -92,7 +92,7 @@ func Start(pythonPath string, threshold float64) (*Detector, error) {
|
||||
}
|
||||
|
||||
// 模型目录:程序运行目录下的 lib(和 application.properties / media-info.db 同一根目录)。
|
||||
// 服务脚本优先加载 lib/320n.onnx;本地没有会自己下载到那里,下载失败再回退用 pip 自带模型。
|
||||
// 服务脚本优先加载 lib/nwfs/320n.onnx;本地没有会自己下载到那里,下载失败再回退用 pip 自带模型。
|
||||
modelDir := ""
|
||||
if wd, err := os.Getwd(); err == nil {
|
||||
modelDir = filepath.Join(wd, "lib")
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user