59 lines
1.8 KiB
Go
59 lines
1.8 KiB
Go
package service
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"strings"
|
|
)
|
|
|
|
type FaceGalleryBuilder struct {
|
|
PythonExe string // path to python.exe
|
|
BuildScript string // path to build_gallery.py
|
|
DetModel string // path to face detection onnx
|
|
RecogModel string // path to face recognition onnx
|
|
DetOutputsConfig string // detection output config JSON
|
|
DatasetDir string // dataset root: person_name/*.jpg
|
|
OutputDB string // output face_gallery.db path
|
|
ExpectedDim int // embedding dimension, default 512
|
|
}
|
|
|
|
func (b *FaceGalleryBuilder) Build() (string, error) {
|
|
if b.ExpectedDim <= 0 {
|
|
b.ExpectedDim = 512
|
|
}
|
|
// Allow python to be just "python" for PATH lookup
|
|
pythonExe := b.PythonExe
|
|
if _, err := os.Stat(b.PythonExe); err != nil {
|
|
pythonExe = "python"
|
|
}
|
|
|
|
// Ensure dataset and output directories exist
|
|
if _, err := os.Stat(b.DatasetDir); os.IsNotExist(err) {
|
|
return "", fmt.Errorf("dataset directory not found: %s", b.DatasetDir)
|
|
}
|
|
outDir := filepath.Dir(b.OutputDB)
|
|
if err := os.MkdirAll(outDir, 0o755); err != nil {
|
|
return "", fmt.Errorf("create output directory: %w", err)
|
|
}
|
|
|
|
args := []string{
|
|
b.BuildScript,
|
|
"--dataset", b.DatasetDir,
|
|
"--db_out", b.OutputDB,
|
|
"--det_model", b.DetModel,
|
|
"--recog_model", b.RecogModel,
|
|
"--det_outputs_config", b.DetOutputsConfig,
|
|
fmt.Sprintf("--expected_dim=%d", b.ExpectedDim),
|
|
}
|
|
|
|
cmd := exec.Command(pythonExe, args...)
|
|
output, err := cmd.CombinedOutput()
|
|
outputStr := strings.TrimSpace(string(output))
|
|
if err != nil {
|
|
return outputStr, fmt.Errorf("build failed: %w\n%s", err, outputStr)
|
|
}
|
|
return outputStr, nil
|
|
}
|