package callback import ( "bytes" "encoding/json" "errors" "fmt" "image" "image/color" _ "image/jpeg" _ "image/png" "io" "math/rand" "os" "path/filepath" "strings" "time" "lone-services/pkg/facepp" "lone-services/pkg/utils" "github.com/disintegration/imaging" ) const defaultExpandRatio = 0.15 func expandImageToTemp(srcPath string, ratio float64) (string, error) { if ratio <= 0 { ratio = defaultExpandRatio } img, err := imaging.Open(srcPath) if err != nil { return utils.StringEmpty, err } img = imaging.Rotate90(img) w := img.Bounds().Dx() h := img.Bounds().Dy() newW := int(float64(w) * (1 + ratio)) newH := int(float64(h) * (1 + ratio)) canvas := imaging.New(newW, newH, color.Black) out := imaging.PasteCenter(canvas, img) ext := strings.ToLower(filepath.Ext(srcPath)) if ext == utils.StringEmpty { ext = ".jpg" } tmpFile, err := os.CreateTemp("", "expand-*"+ext) if err != nil { return utils.StringEmpty, err } tmpPath := tmpFile.Name() if err = tmpFile.Close(); err != nil { _ = os.Remove(tmpPath) return utils.StringEmpty, err } if err = imaging.Save(out, tmpPath); err != nil { _ = os.Remove(tmpPath) return utils.StringEmpty, err } return tmpPath, nil } func saveFaceppTrainImage(srcPath string) (string, error) { saveDir := utils.GetConfigString("facepp.train_image_path") if err := os.MkdirAll(saveDir, os.ModePerm); err != nil { return utils.StringEmpty, err } filename := fmt.Sprintf("%d_%d.jpg", time.Now().UnixMilli(), rand.Intn(1000000)) dstPath := filepath.Join(saveDir, filename) if err := copyLocalFile(srcPath, dstPath); err != nil { return utils.StringEmpty, err } return dstPath, nil } func saveFaceppTrainAnnotations(trainPath string, result json.RawMessage) (string, string, error) { if trainPath == utils.StringEmpty { return utils.StringEmpty, utils.StringEmpty, errors.New("train image path is empty") } base := strings.TrimSuffix(filepath.Base(trainPath), filepath.Ext(trainPath)) origPath, webPath, err := saveFaceppJSON(base, utils.StringEmpty, result) if err != nil { return utils.StringEmpty, utils.StringEmpty, err } labelDir := utils.GetConfigString("facepp.train_label_path") if labelDir == utils.StringEmpty { labelDir = "/data/facepp/train/labels" } if err = os.MkdirAll(labelDir, os.ModePerm); err != nil { return origPath, webPath, err } imgW, imgH, err := imageSize(trainPath) if err != nil { return origPath, webPath, err } yolo, err := facepp.ResultToYOLO(result, imgW, imgH) if err != nil { return origPath, webPath, err } if err = os.WriteFile(filepath.Join(labelDir, base+".txt"), []byte(yolo), 0644); err != nil { return origPath, webPath, err } return origPath, webPath, nil } func saveFaceppJSON(base, suffix string, raw json.RawMessage) (string, string, error) { if len(raw) == 0 || string(raw) == "null" { return utils.StringEmpty, utils.StringEmpty, nil } jsonDir := utils.GetConfigString("facepp.train_json_path") if jsonDir == utils.StringEmpty { jsonDir = "/data/facepp/train/json" } if err := os.MkdirAll(jsonDir, os.ModePerm); err != nil { return utils.StringEmpty, utils.StringEmpty, err } origPath := filepath.Join(jsonDir, base+suffix+".json") if err := writePrettyJSON(origPath, raw); err != nil { return utils.StringEmpty, utils.StringEmpty, err } return origPath, filepath.Base(origPath), nil } func writePrettyJSON(path string, raw json.RawMessage) error { if len(raw) == 0 || string(raw) == "null" { raw = json.RawMessage("{}") } var buf bytes.Buffer if err := json.Indent(&buf, raw, "", " "); err != nil { buf.Reset() buf.Write(raw) } return os.WriteFile(path, buf.Bytes(), 0644) } func imageSize(path string) (int, int, error) { f, err := os.Open(path) if err != nil { return 0, 0, err } defer f.Close() cfg, _, err := image.DecodeConfig(f) if err != nil { return 0, 0, err } return cfg.Width, cfg.Height, nil } func copyLocalFile(src, dst string) error { in, err := os.Open(src) if err != nil { return err } defer in.Close() out, err := os.Create(dst) if err != nil { return err } defer out.Close() _, err = io.Copy(out, in) return err }