feat: 色情识别模型本地化(lib/320n.onnx),缺失时自动下载并回退自带模型

This commit is contained in:
2025-06-11 21:15:24 +08:00
parent d4f9dc362e
commit 01a268663f
5 changed files with 85 additions and 6 deletions
+10 -3
View File
@@ -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 {