Files
mediadevices/examples/facedetection/detector.go
2020-03-25 23:46:33 -04:00

119 lines
2.5 KiB
Go

package main
import (
"image"
"image/color"
"image/draw"
"io/ioutil"
"log"
"github.com/disintegration/imaging"
pigo "github.com/esimov/pigo/core"
)
var (
cascade []byte
err error
classifier *pigo.Pigo
)
func imgToGrayscale(img image.Image) []uint8 {
bounds := img.Bounds()
flatten := bounds.Dy() * bounds.Dx()
grayImg := make([]uint8, flatten)
i := 0
for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
for x := bounds.Min.X; x < bounds.Max.X; x++ {
pix := img.At(x, y)
grayPix := color.GrayModel.Convert(pix).(color.Gray)
grayImg[i] = grayPix.Y
i++
}
}
return grayImg
}
// clusterDetection runs Pigo face detector core methods
// and returns a cluster with the detected faces coordinates.
func clusterDetection(img image.Image) []pigo.Detection {
grayscale := imgToGrayscale(img)
bounds := img.Bounds()
cParams := pigo.CascadeParams{
MinSize: 100,
MaxSize: 600,
ShiftFactor: 0.15,
ScaleFactor: 1.1,
ImageParams: pigo.ImageParams{
Pixels: grayscale,
Rows: bounds.Dy(),
Cols: bounds.Dx(),
Dim: bounds.Dx(),
},
}
if len(cascade) == 0 {
cascade, err = ioutil.ReadFile("facefinder")
if err != nil {
log.Fatalf("Error reading the cascade file: %s", err)
}
p := pigo.NewPigo()
// Unpack the binary file. This will return the number of cascade trees,
// the tree depth, the threshold and the prediction from tree's leaf nodes.
classifier, err = p.Unpack(cascade)
if err != nil {
log.Fatalf("Error unpacking the cascade file: %s", err)
}
}
// Run the classifier over the obtained leaf nodes and return the detection results.
// The result contains quadruplets representing the row, column, scale and detection score.
dets := classifier.RunCascade(cParams, 0.0)
// Calculate the intersection over union (IoU) of two clusters.
dets = classifier.ClusterDetections(dets, 0)
return dets
}
func drawCircle(img draw.Image, x0, y0, r int, c color.Color) {
x, y, dx, dy := r-1, 0, 1, 1
err := dx - (r * 2)
for x > y {
img.Set(x0+x, y0+y, c)
img.Set(x0+y, y0+x, c)
img.Set(x0-y, y0+x, c)
img.Set(x0-x, y0+y, c)
img.Set(x0-x, y0-y, c)
img.Set(x0-y, y0-x, c)
img.Set(x0+y, y0-x, c)
img.Set(x0+x, y0-y, c)
if err <= 0 {
y++
err += dy
dy += 2
}
if err > 0 {
x--
dx += 2
err += dx - (r * 2)
}
}
}
func markFaces(img image.Image) image.Image {
nrgba := imaging.Clone(img)
dets := clusterDetection(img)
for _, det := range dets {
if det.Q < 5.0 {
continue
}
drawCircle(nrgba, det.Col, det.Row, det.Scale/2, color.Black)
}
return nrgba
}