feat: add facepp

pkg/facepp/yolo.go

@@ -0,0 +1,127 @@
+package facepp
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+const (
+	YOLOClassAcne      = 0
+	YOLOClassMole      = 1
+	YOLOClassBrownSpot = 2
+)
+
+type yoloField struct {
+	Key     string
+	ClassID int
+}
+
+var yoloRectangleFields = []yoloField{
+	{Key: "acne_pustule", ClassID: YOLOClassAcne},
+	{Key: "closed_comedones", ClassID: YOLOClassAcne},
+	{Key: "acne_mark", ClassID: YOLOClassAcne},
+	{Key: "acne", ClassID: YOLOClassAcne},
+	{Key: "acne_nodule", ClassID: YOLOClassAcne},
+	{Key: "mole", ClassID: YOLOClassMole},
+	{Key: "brown_spot", ClassID: YOLOClassBrownSpot},
+}
+
+type yoloBox struct {
+	ClassID int
+	XC      float64
+	YC      float64
+	W       float64
+	H       float64
+}
+
+type yoloRect struct {
+	Top    float64 `json:"top"`
+	Left   float64 `json:"left"`
+	Width  float64 `json:"width"`
+	Height float64 `json:"height"`
+}
+
+type yoloLocated struct {
+	Rectangle []yoloRect `json:"rectangle"`
+	Rect      []yoloRect `json:"rect"`
+}
+
+func ResultToYOLO(result json.RawMessage, imgW, imgH int) (string, error) {
+	if len(result) == 0 || string(result) == "null" {
+		return "", nil
+	}
+	if imgW <= 0 || imgH <= 0 {
+		return "", fmt.Errorf("facepp: invalid image size %dx%d", imgW, imgH)
+	}
+
+	var root map[string]json.RawMessage
+	if err := json.Unmarshal(result, &root); err != nil {
+		return "", err
+	}
+
+	imgWf := float64(imgW)
+	imgHf := float64(imgH)
+	var boxes []yoloBox
+	for _, field := range yoloRectangleFields {
+		raw, ok := root[field.Key]
+		if !ok {
+			continue
+		}
+		for _, rect := range locatedRectangles(raw) {
+			box, ok := pixelToYOLO(field.ClassID, rect.Left, rect.Top, rect.Width, rect.Height, imgWf, imgHf)
+			if !ok {
+				continue
+			}
+			boxes = append(boxes, box)
+		}
+	}
+
+	var b strings.Builder
+	for _, box := range boxes {
+		fmt.Fprintf(&b, "%d %.6f %.6f %.6f %.6f\n", box.ClassID, box.XC, box.YC, box.W, box.H)
+	}
+	return b.String(), nil
+}
+
+func locatedRectangles(raw json.RawMessage) []yoloRect {
+	var obj yoloLocated
+	if err := json.Unmarshal(raw, &obj); err != nil {
+		return nil
+	}
+	if len(obj.Rectangle) > 0 {
+		return obj.Rectangle
+	}
+	return obj.Rect
+}
+
+func pixelToYOLO(classID int, xMin, yMin, boxW, boxH, imgW, imgH float64) (yoloBox, bool) {
+	xMax := xMin + boxW
+	yMax := yMin + boxH
+	if xMin < 0 {
+		xMin = 0
+	}
+	if yMin < 0 {
+		yMin = 0
+	}
+	if xMax > imgW {
+		xMax = imgW
+	}
+	if yMax > imgH {
+		yMax = imgH
+	}
+
+	boxW = xMax - xMin
+	boxH = yMax - yMin
+	if boxW <= 0 || boxH <= 0 {
+		return yoloBox{}, false
+	}
+
+	return yoloBox{
+		ClassID: classID,
+		XC:      (xMin + boxW/2) / imgW,
+		YC:      (yMin + boxH/2) / imgH,
+		W:       boxW / imgW,
+		H:       boxH / imgH,
+	}, true
+}