cctv/index.go

package main

import (
	"image/color"
	"io"
	"log"
	"os"
	"time"

	pigo "github.com/esimov/pigo/core"
	"github.com/fogleman/gg"
	"github.com/unixpickle/ffmpego"
)

var f *os.File
var dc *gg.Context

func videoFileFacetest(withOutput bool) {
	cascadeFile, err := os.ReadFile("./cascade/facefinder")
	if err != nil {
		log.Fatalf("Error reading the cascade file: %v", err)
	}

	vr, _ := ffmpego.NewVideoReaderResampled("face2.mp4", 30)

	e := os.Remove("cropped.mp4")
	if os.IsNotExist(e) {
		// handle the case where the file doesn't exist
	}
	vf, _ := ffmpego.NewVideoWriter("cropped.mp4", vr.VideoInfo().Width, vr.VideoInfo().Height, 1)

	var vw *ffmpego.VideoWriter
	if withOutput {
		e := os.Remove("output.mp4")
		if os.IsNotExist(e) {
			// handle the case where the file doesn't exist
		}
		vw, _ = ffmpego.NewVideoWriter("output.mp4", vr.VideoInfo().Width, vr.VideoInfo().Height, 30)
	}
	i := 0
	start := time.Now()
	for {
		i++
		frame, err := vr.ReadFrame()
		if err == io.EOF {
			break
		}

		pixels := pigo.RgbToGrayscale(frame)

		cols, rows := frame.Bounds().Max.X, frame.Bounds().Max.Y

		cParams := pigo.CascadeParams{
			MinSize:     rows / 18,
			MaxSize:     rows,
			ShiftFactor: 0.1,
			ScaleFactor: 1.5,

			ImageParams: pigo.ImageParams{
				Pixels: pixels,
				Rows:   rows,
				Cols:   cols,
				Dim:    cols,
			},
		}
		pigoFunc := pigo.NewPigo()
		classifier, err := pigoFunc.Unpack(cascadeFile)
		if err != nil {
			log.Fatalf("Error reading the cascade file: %s", err)
		}
		angle := 0.0 // cascade rotation angle. 0.0 is 0 radians and 1.0 is 2*pi radians

		// 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, angle)
		// // Calculate the intersection over union (IoU) of two clusters.

		// print(len(dets))
		// print("_")
		// println(len(dets2))

		var qTresh float32
		qTresh = 6.5
		var qTresh2 float32
		qTresh2 = 10.0
		goodQ := []pigo.Detection{}
		for i := range dets {
			if dets[i].Q > qTresh {
				goodQ = append(goodQ, dets[i])
			}
		}
		if len(goodQ) > 0 {
			dets2 := classifier.ClusterDetections(goodQ, 0.1)
			if len(dets2) > 0 {
				foundFace := false
				if withOutput {
					dc = gg.NewContext(cols, rows)
					dc.DrawImage(frame, 0, 0)
				}
				for i := 0; i < len(dets2); i++ {
					if dets2[i].Q > qTresh2 {
						println("we got a winner here")
						foundFace = true
						//do somethng here
						//send cropped face to server
						// rgba := frame.(*image.RGBA)

						// cropped := rgba.SubImage(image.Rect(dets2[i].Col-dets2[i].Scale/2, dets2[i].Row-dets2[i].Scale/2, dets2[i].Scale, dets2[i].Scale))
						// file, _ := os.Create(fmt.Sprintf(" %d.jpg", i))
						// png.Encode(file, cropped)
						// cropped := imaging.Crop(frame, image.Rect(dets2[i].Col-dets2[i].Scale/2, dets2[i].Row-dets2[i].Scale/2, dets2[i].Scale, dets2[i].Scale))
						// dc = gg.NewContext(cropped.Rect.Dx(), cropped.Rect.Dy())
						// dc.DrawImage(cropped, 0, 0)
						// vf.WriteFrame(dc.Image())
						dc.SetMask(dc.AsMask())
						if withOutput {
							dc.DrawRectangle(
								float64(dets2[i].Col-dets2[i].Scale/2),
								float64(dets2[i].Row-dets2[i].Scale/2),
								float64(dets2[i].Scale),
								float64(dets2[i].Scale),
							)
							dc.SetLineWidth(float64(rows / 43))
							dc.SetStrokeStyle(gg.NewSolidPattern(color.RGBA{R: 255, G: 0, B: 0, A: 255}))
							dc.Stroke()
							dc.InvertMask()
						}
					}
				}
				if withOutput && foundFace {
					vw.WriteFrame(dc.Image())
				}
			}
		}
		log.Printf("Loop took %s", time.Since(start))
	}
	elapsed := time.Since(start)
	log.Printf("Loop took %s", elapsed)
	vr.Close()
	vf.Close()

	if withOutput {
		vw.Close()
	}

	return
}

func main() {
	videoFileFacetest(true)
}