in-process video machine vision alg

frame_reader.go

package main

import (
	"fmt"
	"image"
	"io"
	"log"
	"os"

	"github.com/nareix/joy4/av/avutil"
	"github.com/nareix/joy4/cgo/ffmpeg"
	"github.com/nareix/joy4/format"
)

func init() {
	format.RegisterAll()
}

func main() {
	if len(os.Args) < 2 {
		log.Fatalln("Provide a path to a video file to analyze.")
	}

	file, err := avutil.Open(os.Args[1])
	if err != nil {
		log.Fatal(err)
	}

	streams, _ := file.Streams()
	var vdecoder *ffmpeg.VideoDecoder
	var streamIdx int8
	for idx, stream := range streams {
		if stream.Type().IsVideo() {
			vdecoder, err = ffmpeg.NewVideoDecoder(stream)
			if err != nil {
				log.Fatal(err)
			}
			streamIdx = int8(idx)
		}
	}
	if vdecoder == nil {
		log.Fatalln("Didn't find a video stream in the file")
	}

	pktC := make(chan []byte, 5)
	go func() {
		for {
			if pkt, err := file.ReadPacket(); err != nil {
				if err == io.EOF {
					break
				}
				log.Fatal(err)
			} else if pkt.Idx == streamIdx {
				pktC <- pkt.Data
			}
		}
		file.Close()
		close(pktC)
	}()

	for pktData := range pktC {
		img, err := vdecoder.Decode(pktData)
		if err != nil {
			log.Fatal(err)
		}
		if img != nil {
			handleFrame(img)
		}
	}
	// TODO: flush out any remaining frames
	// the binding library doesn't support sending zero-byte frames in
	// close(frameC)
}

var frameIdx int = 0

func handleFrame(frame *ffmpeg.VideoFrame) {
	frameIdx++
	fmt.Printf("{\"frame\":%d,\"hasMarker\":%v}\n", frameIdx, containsSyncMarker(frame.Image))
	frame.Free()
}

func containsSyncMarker(img image.YCbCr) bool {
	var width = img.Bounds().Dx()
	var height = img.Bounds().Dy()
	var widthToCheck = int(float32(width) * 0.08)
	var heightToCheck = int(float32(height) * 0.08)

	for x := 0; x < widthToCheck; x++ {
		for y := height - heightToCheck; y < height; y++ {
			r, g, b, _ := img.At(x, y).RGBA()
			// log.Println(r>>8, g>>8, b>>8)
			if r>>8 > 10 || g>>8 < 240 || b>>8 > 10 {
				return false
			}
		}
	}

	return true
}