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

package callback

import (
	"encoding/json"
	"errors"
	"fmt"
	"image/color"
	"math/rand"
	"os"
	"path/filepath"
	"strings"
	"time"

	"lone-services/pkg/facepp"
	"lone-services/pkg/modelbase"
	pkgoss "lone-services/pkg/oss"
	"lone-services/pkg/utils"
	"lone-services/services/analysis/internal/dao"
	"lone-services/services/analysis/internal/model"

	"github.com/zeromicro/go-zero/core/logx"
)

var faceppMapStroke = map[string]color.RGBA{
	dao.FaceppMapTypeAcne:            {R: 0xFE, G: 0x01, B: 0x00, A: 0xFF},
	dao.FaceppMapTypeClosedComedones: {R: 0x00, G: 0xFF, B: 0x00, A: 0xFF},
	dao.FaceppMapTypeBrownSpot:       {R: 0x7E, G: 0x2A, B: 0x28, A: 0xFF},
	dao.FaceppMapTypeMole:            {R: 0x00, G: 0x66, B: 0xFF, A: 0xFF},
}

func RunFacepp(taskID, finalDir string) {
	defer func() {
		if r := recover(); r != nil {
			logx.Errorf("facepp panic task_id=%s: %v", taskID, r)
		}
	}()

	originalPath := filepath.Join(finalDir, "cross_mid.jpg")
	if _, err := os.Stat(originalPath); err != nil {
		if os.IsNotExist(err) {
			logx.Infof("skip facepp: cross_mid.jpg not found in %s", finalDir)
			return
		}
		logx.Errorf("stat cross_mid.jpg failed: %v", err)
		return
	}

	expandPath, err := expandImageToTemp(originalPath, defaultExpandRatio)
	if err != nil {
		logx.Errorf("expand image failed: %v", err)
		return
	}
	defer os.Remove(expandPath)

	trainPath, err := saveFaceppTrainImage(expandPath)
	if err != nil {
		logx.Errorf("save facepp train image failed: %v", err)
		return
	}

	client := facepp.NewClient(
		utils.GetConfigString("facepp.api_key"),
		utils.GetConfigString("facepp.api_secret"),
	)
	resp, analyzeErr := client.SkinAnalyzeFromFile(expandPath, facepp.SkinAnalyzeOptions{
		ReturnMarks: facepp.DefaultReturnMarks,
	})

	resultPath := utils.StringEmpty
	resultWeb := utils.StringEmpty
	leftWeb := utils.StringEmpty
	rightWeb := utils.StringEmpty

	if analyzeErr != nil {
		logx.Errorf("facepp skin analyze failed: %v", analyzeErr)
	} else if resp != nil {
		var annErr error
		resultPath, resultWeb, annErr = saveFaceppTrainAnnotations(trainPath, resp.Result)
		if annErr != nil {
			logx.Errorf("save facepp train annotations failed: %v", annErr)
		}

		base := strings.TrimSuffix(filepath.Base(trainPath), filepath.Ext(trainPath))
		_, leftWeb, err = saveFaceppJSON(base, "_left", resp.LeftSideResult)
		if err != nil {
			logx.Errorf("save facepp left json failed: %v", err)
		}
		_, rightWeb, err = saveFaceppJSON(base, "_right", resp.RightSideResult)
		if err != nil {
			logx.Errorf("save facepp right json failed: %v", err)
		}

		if mapErr := saveFaceppMaps(taskID, expandPath, resp.Result); mapErr != nil {
			logx.Errorf("save facepp maps failed: %v", mapErr)
		}
	}

	if err = saveFaceppRecord(taskID, originalPath, trainPath, resultPath, resultWeb, leftWeb, rightWeb, resp, analyzeErr); err != nil {
		logx.Errorf("save facepp record failed: %v", err)
	}
}

func saveFaceppRecord(taskID, originalPath, expandPath, resultPath, resultWeb, leftWeb, rightWeb string, resp *facepp.SkinAnalyzeResponse, analyzeErr error) error {
	ossResultPath := utils.StringEmpty
	if resultPath != utils.StringEmpty {
		uploaded, err := uploadJSONToOSS(resultPath)
		if err != nil {
			logx.Errorf("upload facepp json to oss failed: %v", err)
		} else {
			ossResultPath = uploaded
		}
	}

	record := dao.FaceppAnalysis{
		TaskId:          taskID,
		OriginalPath:    originalPath,
		ExpandPath:      expandPath,
		Result:          resultWeb,
		LeftSideResult:  leftWeb,
		RightSideResult: rightWeb,
		ResultPath:      resultPath,
		ResultOssPath:   ossResultPath,
		CreatedAt:       utils.Now(),
	}
	if resp != nil {
		record.RequestId = resp.RequestID
		record.FaceTop = resp.FaceRectangle.Top
		record.FaceLeft = resp.FaceRectangle.Left
		record.FaceWidth = resp.FaceRectangle.Width
		record.FaceHeight = resp.FaceRectangle.Height
		record.TimeUsed = resp.TimeUsed
		record.ErrorMessage = resp.ErrorMessage
	}
	if record.ErrorMessage == utils.StringEmpty && analyzeErr != nil {
		record.ErrorMessage = analyzeErr.Error()
	}

	m := model.FaceppAnalysisModel{}.Init()
	return m.Create(&record)
}

func uploadJSONToOSS(localPath string) (string, error) {
	data, err := os.ReadFile(localPath)
	if err != nil {
		return utils.StringEmpty, err
	}
	fileDir := utils.GetConfigString("oss.fileDir")
	fileName := fmt.Sprintf("%d%d.json", time.Now().UnixNano(), rand.Intn(10000))
	res, err := pkgoss.Upload(pkgoss.UploadInput{
		Filename:    fileName,
		Content:     data,
		ContentType: "application/json",
		Dir:         fileDir,
	})
	if err != nil {
		return utils.StringEmpty, err
	}
	return res.URL, nil
}

func saveFaceppMaps(taskID, imagePath string, result json.RawMessage) error {
	if imagePath == utils.StringEmpty {
		return errors.New("facepp map image path empty")
	}
	now := utils.Now()
	for _, meta := range dao.FaceppMapTypeMeta {
		keys := facepp.MapTypeKeys[meta.Type]
		pixelBoxes, err := facepp.ExtractPixelBoxes(result, keys...)
		if err != nil {
			return fmt.Errorf("extract %s boxes: %w", meta.Type, err)
		}

		tmpFile, err := os.CreateTemp("", "facepp-map-*.jpg")
		if err != nil {
			return err
		}
		tmpPath := tmpFile.Name()
		_ = tmpFile.Close()

		stroke := faceppMapStroke[meta.Type]
		imgW, imgH, err := facepp.DrawBoxesOnFile(imagePath, tmpPath, pixelBoxes, stroke, 3)
		if err != nil {
			_ = os.Remove(tmpPath)
			return fmt.Errorf("draw %s: %w", meta.Type, err)
		}

		ossURL, uploadErr := uploadFaceppImageToOSS(tmpPath)
		_ = os.Remove(tmpPath)
		if uploadErr != nil {
			return fmt.Errorf("upload %s: %w", meta.Type, uploadErr)
		}

		if err = upsertFaceppMapImage(dao.FaceppMap{
			TaskId:      taskID,
			Type:        meta.Type,
			Path:        ossURL,
			ImageWidth:  imgW,
			ImageHeight: imgH,
			CreatedAt:   now,
			UpdatedAt:   now,
		}); err != nil {
			return err
		}
	}
	return nil
}

func upsertFaceppMapImage(record dao.FaceppMap) error {
	m := model.FaceppMapModel{}.Init()
	var existing dao.FaceppMap
	w := modelbase.Params{
		Eq: map[string]string{
			"task_id": record.TaskId,
			"type":    record.Type,
		},
	}
	if err := m.GetOne(w, &existing); err != nil {
		return err
	}
	if existing.Id > utils.NumberZero {
		_, err := m.Edit(w, map[string]interface{}{
			"path":         record.Path,
			"image_width":  record.ImageWidth,
			"image_height": record.ImageHeight,
			"updated_at":   record.UpdatedAt,
		})
		return err
	}
	return m.Create(&record)
}

func uploadFaceppImageToOSS(localPath string) (string, error) {
	data, err := os.ReadFile(localPath)
	if err != nil {
		return utils.StringEmpty, err
	}
	imageDir := utils.GetConfigString("oss.imageDir")
	fileName := fmt.Sprintf("%d%d.jpg", time.Now().UnixNano(), rand.Intn(10000))
	res, err := pkgoss.Upload(pkgoss.UploadInput{
		Filename:    fileName,
		Content:     data,
		ContentType: "image/jpeg",
		Dir:         imageDir,
	})
	if err != nil {
		return utils.StringEmpty, err
	}
	return res.URL, nil
}