lone-services/services/analysis/internal/callback/facepp_file.go

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
}