upd
This commit is contained in:
@@ -60,6 +60,7 @@ const requestTimeout = 30 * time.Second
|
||||
type Detector struct {
|
||||
url string // 本地服务地址,如 http://127.0.0.1:54321
|
||||
threshold float64 // 判定阈值 0~1
|
||||
workers int // Python 推理 worker 数量
|
||||
client *http.Client
|
||||
cmd *exec.Cmd // python 进程句柄(Stop 时杀掉)
|
||||
scriptDir string // 临时脚本所在目录(Stop 时清理)
|
||||
@@ -75,12 +76,16 @@ type Detector struct {
|
||||
// Start 拉起本地 python 检测服务并等它就绪。
|
||||
// - pythonPath:python 可执行文件路径;为空时在 PATH 里找 python/python3。
|
||||
// - threshold:判定阈值 0~1。
|
||||
// - workers:Python 推理 worker 数量;<=0 表示使用默认值(4)。
|
||||
//
|
||||
// 找不到 python、装不了 nudenet、或服务起不来时返回错误(调用方跳过识别即可)。
|
||||
func Start(pythonPath string, threshold float64) (*Detector, error) {
|
||||
func Start(pythonPath string, threshold float64, workers int) (*Detector, error) {
|
||||
if threshold <= 0 || threshold >= 1 {
|
||||
threshold = 0.5 // 阈值传得不合理就用默认
|
||||
}
|
||||
if workers <= 0 {
|
||||
workers = 4 // 默认 4 个 worker
|
||||
}
|
||||
// 找 python
|
||||
py := resolvePython(pythonPath)
|
||||
if py == "" {
|
||||
@@ -107,6 +112,7 @@ func Start(pythonPath string, threshold float64) (*Detector, error) {
|
||||
|
||||
d := &Detector{
|
||||
threshold: threshold,
|
||||
workers: workers,
|
||||
client: &http.Client{Timeout: requestTimeout},
|
||||
python: py,
|
||||
modelDir: modelDir,
|
||||
@@ -135,6 +141,8 @@ func (d *Detector) launch() error {
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cmd := exec.CommandContext(ctx, d.python, filepath.Join(d.scriptDir, "server.py"), fmt.Sprintf("%d", port), d.modelDir)
|
||||
// 设置环境变量,传递 worker 数量给 Python 服务
|
||||
cmd.Env = append(os.Environ(), fmt.Sprintf("KM_NSFW_WORKERS=%d", d.workers))
|
||||
// 标准输出/错误都接到程序自身的输出:
|
||||
// - 扫描期间 os.Stdout 会被接管进"运行日志",这样 python 的报错(崩溃原因)也能在前端看到
|
||||
// - 方便排查服务为什么挂掉
|
||||
|
||||
Reference in New Issue
Block a user