Files
mediadevices/examples/facedetection/main.go
Lukas Herman 18b81bfba6 WIP
2020-11-04 05:43:40 +00:00

118 lines
2.8 KiB
Go

package main
import (
"image"
"io/ioutil"
"log"
"time"
pigo "github.com/esimov/pigo/core"
"github.com/pion/mediadevices"
_ "github.com/pion/mediadevices/pkg/driver/camera" // This is required to register camera adapter
"github.com/pion/mediadevices/pkg/frame"
"github.com/pion/mediadevices/pkg/prop"
)
const (
confidenceLevel = 5.0
)
var (
cascade []byte
classifier *pigo.Pigo
)
func must(err error) {
if err != nil {
panic(err)
}
}
func detectFace(frame *image.YCbCr) bool {
bounds := frame.Bounds()
cascadeParams := pigo.CascadeParams{
MinSize: 100,
MaxSize: 600,
ShiftFactor: 0.15,
ScaleFactor: 1.1,
ImageParams: pigo.ImageParams{
Pixels: frame.Y, // Y in YCbCr should be enough to detect faces
Rows: bounds.Dy(),
Cols: bounds.Dx(),
Dim: bounds.Dx(),
},
}
// 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(cascadeParams, 0.0)
// Calculate the intersection over union (IoU) of two clusters.
dets = classifier.ClusterDetections(dets, 0)
for _, det := range dets {
if det.Q >= confidenceLevel {
return true
}
}
return false
}
func main() {
// prepare face detector
var err error
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)
}
devices := mediadevices.EnumerateDevices()
deviceID := ""
for _, device := range devices {
if device.Label == "video0" {
deviceID = device.DeviceID
}
}
mediaStream, err := mediadevices.GetUserMedia(mediadevices.MediaStreamConstraints{
Video: func(c *mediadevices.MediaTrackConstraints) {
c.DeviceID = prop.StringExact(deviceID)
c.FrameFormat = prop.FrameFormatExact(frame.FormatUYVY)
c.Width = prop.Int(640)
c.Height = prop.Int(480)
},
})
must(err)
// since we're trying to access the raw data, we need to cast Track to its real type, *mediadevices.VideoTrack
videoTrack := mediaStream.GetVideoTracks()[0].(*mediadevices.VideoTrack)
defer videoTrack.Close()
videoReader := videoTrack.NewReader(false)
// To save resources, we can simply use 4 fps to detect faces.
ticker := time.NewTicker(time.Millisecond * 250)
defer ticker.Stop()
for range ticker.C {
frame, release, err := videoReader.Read()
must(err)
// Since we asked the frame format to be exactly YUY2 in GetUserMedia, we can guarantee that it must be YCbCr
if detectFace(frame.(*image.YCbCr)) {
log.Println("Detect a face")
}
release()
}
}