detect face

who-face.go

package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"regexp"

	"github.com/Kagami/go-face"
)

// TODO: add description and MV to github repo

const (
	dataDir   = "../go-face-testdata"
	modelsDir = dataDir + "/models"
	imagesDir = dataDir + "/images_2"
	testDir   = dataDir + "/images-test"
)

var testFace = flag.String("f", "testdata_2.jpg", "Which face will try to find in filename")
var rgx = regexp.MustCompile(`(\d+)?.(jpg|jpeg)`)

// TrainData c
type TrainData struct {
	labels      []string
	catIDs      []int32
	descriptors []face.Descriptor
}

// Face with descriptor to each label
type Face struct {
	Descr face.Descriptor
	Label string
}

func init() {
	flag.Parse()
}

func main() {
	rec, err := face.NewRecognizer(modelsDir)
	if err != nil {
		log.Fatal(err)
	}
	defer rec.Close()

	// or faceData() and trainData()
	faces := getFaceData(rec)
	tdata := getTrainData(faces)
	rec.SetSamples(tdata.descriptors, tdata.catIDs)

	whoFaces, err := rec.RecognizeFile(testDir + "/" + *testFace)
	if err != nil {
		log.Fatal(err)
	}
	for _, sample := range whoFaces {
		id := rec.Classify(sample.Descriptor)
		label := tdata.labels[id]
		fmt.Printf("for %v it was %v?\n", *testFace, label)
	}
}

func getFaceData(rec *face.Recognizer) []*Face {
	files, err := ioutil.ReadDir(imagesDir)
	if err != nil {
		log.Fatal(err)
	}

	var faces []*Face

	for _, f := range files {
		filename := f.Name()
		name := rgx.ReplaceAllString(filename, "")
		sample, err := rec.RecognizeSingleFile(imagesDir + "/" + filename)
		if err != nil {
			log.Fatalf("Can't recognize: %v", err)
		}
		if sample == nil {
			log.Printf("Not a single face on the image: %v", filename)
			continue
		}

		face := &Face{Descr: sample.Descriptor, Label: name}
		faces = append(faces, face)
	}

	return faces
}

func getTrainData(data []*Face) *TrainData {
	var prevLabel string
	var catID int32 = -1
	var labels []string
	var catIDs = make([]int32, len(data))
	var descriptors = make([]face.Descriptor, len(data))

	for i, d := range data {
		if d.Label != prevLabel {
			catID++
			labels = append(labels, d.Label)
		}
		catIDs[i] = catID
		descriptors[i] = d.Descr
		prevLabel = d.Label
	}

	tdata := &TrainData{
		labels,
		catIDs,
		descriptors,
	}
	return tdata
}