feat: add facepp

services/analysis/internal/callback/facepp_file.go

@@ -0,0 +1,168 @@
+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 saveDir == utils.StringEmpty {
+		saveDir = "/data/facepp/train/images"
+	}
+	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
+}