lone-services/pkg/facepp/yolo.go

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
}