lone-services/pkg/facepp/annotate.go

package facepp

import (
	"image"
	"image/color"
	"image/draw"
	"image/jpeg"
	"os"

	"github.com/disintegration/imaging"
)

func DrawBoxesOnFile(srcPath, dstPath string, boxes []PixelBox, stroke color.RGBA, thickness int) (int, int, error) {
	img, err := imaging.Open(srcPath, imaging.AutoOrientation(true))
	if err != nil {
		return 0, 0, err
	}
	if thickness <= 0 {
		thickness = 3
	}

	bounds := img.Bounds()
	dst := image.NewNRGBA(bounds)
	draw.Draw(dst, bounds, img, bounds.Min, draw.Src)

	for _, box := range boxes {
		drawRectBorder(dst, box, stroke, thickness)
	}

	if err = saveJPEG(dst, dstPath); err != nil {
		return 0, 0, err
	}
	return bounds.Dx(), bounds.Dy(), nil
}

func drawRectBorder(img *image.NRGBA, box PixelBox, c color.RGBA, thickness int) {
	x0 := int(box.Left)
	y0 := int(box.Top)
	x1 := int(box.Left + box.Width)
	y1 := int(box.Top + box.Height)

	b := img.Bounds()
	if x0 < b.Min.X {
		x0 = b.Min.X
	}
	if y0 < b.Min.Y {
		y0 = b.Min.Y
	}
	if x1 > b.Max.X {
		x1 = b.Max.X
	}
	if y1 > b.Max.Y {
		y1 = b.Max.Y
	}
	if x1 <= x0 || y1 <= y0 {
		return
	}

	for t := 0; t < thickness; t++ {
		for x := x0; x < x1; x++ {
			setNRGBA(img, x, y0+t, c)
			setNRGBA(img, x, y1-1-t, c)
		}
		for y := y0; y < y1; y++ {
			setNRGBA(img, x0+t, y, c)
			setNRGBA(img, x1-1-t, y, c)
		}
	}
}

func setNRGBA(img *image.NRGBA, x, y int, c color.RGBA) {
	b := img.Bounds()
	if x < b.Min.X || x >= b.Max.X || y < b.Min.Y || y >= b.Max.Y {
		return
	}
	img.SetNRGBA(x, y, color.NRGBA{R: c.R, G: c.G, B: c.B, A: 255})
}

func saveJPEG(img image.Image, path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	defer f.Close()
	return jpeg.Encode(f, img, &jpeg.Options{Quality: 92})
}