website/pkg/photo/face_score.go

package photo

import (
	"errors"
	"os"

	"code.nonshy.com/nonshy/website/pkg/config"
	"code.nonshy.com/nonshy/website/pkg/log"
	pigo "github.com/esimov/pigo/core"
)

// Functionality to do with face detection in pictures.

var (
	faceScoreReady bool
	faceClassifier *pigo.Pigo
	cascadeFile    []byte
)

// InitFaceScore initializes the face recognition library (esimov/pigo).
func InitFaceScore() {
	if faceScoreReady {
		return
	}

	if !config.Current.FaceScore.Enabled {
		log.Error("InitFaceScore: not enabled in settings, face detection will not run")
		return
	}

	// Load the cascade file needed for face detection.
	data, err := os.ReadFile(config.Current.FaceScore.CascadeFile)
	if err != nil {
		log.Error("InitFaceScore: could not load cascade file (%s): %s", config.Current.FaceScore.CascadeFile, err)
		return
	}

	cascadeFile = data

	log.Info("Initializing FaceScore with cascade file (%d bytes)", len(cascadeFile))

	faceClassifier = pigo.NewPigo()
	faceClassifier, err = faceClassifier.Unpack(cascadeFile)
	if err != nil {
		log.Error("InitFaceScore: could not unpack the cascade file: %s", err)
		return
	}

	faceScoreReady = true
}

// ComputeFaceScore checks a photo on disk and returns the detected face score.
func ComputeFaceScore(filename string) (float64, error) {
	if !faceScoreReady {
		return 0, errors.New("face detection is not available")
	}

	src, err := pigo.GetImage(filename)
	if err != nil {
		return 0, err
	}

	var (
		pixels     = pigo.RgbToGrayscale(src)
		cols, rows = src.Bounds().Max.X, src.Bounds().Max.Y
		cParams    = pigo.CascadeParams{
			MinSize:     20,
			MaxSize:     1000,
			ShiftFactor: 0.1,
			ScaleFactor: 1.1,
			ImageParams: pigo.ImageParams{
				Pixels: pixels,
				Rows:   rows,
				Cols:   cols,
				Dim:    cols,
			},
		}
	)

	// Run the classifier.
	dets := faceClassifier.RunCascade(cParams, 0.0)
	for _, row := range dets {
		log.Warn("%+v", row)
	}

	// Note: the classifier may return multiple matched faces, return the highest score.
	var highest float32
	for _, row := range dets {
		if row.Q > highest {
			highest = row.Q
		}
	}

	return float64(highest), nil
}