av/filter/mog.go

// +build !circleci

/*
DESCRIPTION
  A filter that detects motion and discards frames without motion. The
  filter uses a Mixture of Gaussians method (MoG) to determine what is
  background and what is foreground.

AUTHORS
  Scott Barnard <scott@ausocean.org>

LICENSE
  mog.go is Copyright (C) 2019 the Australian Ocean Lab (AusOcean)

  It is free software: you can redistribute it and/or modify them
  under the terms of the GNU General Public License as published by the
  Free Software Foundation, either version 3 of the License, or (at your
  option) any later version.

  It is distributed in the hope that it will be useful, but WITHOUT
  ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
  FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
  for more details.

  You should have received a copy of the GNU General Public License
  in gpl.txt.  If not, see http://www.gnu.org/licenses.
*/

package filter

import (
	"fmt"
	"image"
	"image/color"
	"io"

	"gocv.io/x/gocv"
)

// MOGFilter is a filter that provides basic motion detection. MoG is short for
// Mixture of Gaussians method.
type MOGFilter struct {
	dst     io.WriteCloser
	area    float64
	bs      *gocv.BackgroundSubtractorMOG2
	knl     gocv.Mat
	debug   bool
	windows []*gocv.Window
	hold    [][]byte
	hf      int
}

var hfCount int = 0

// NewMOGFilter returns a pointer to a new MOGFilter struct.
func NewMOGFilter(dst io.WriteCloser, area, threshold float64, history int, debug bool, hf int) *MOGFilter {
	bs := gocv.NewBackgroundSubtractorMOG2WithParams(history, threshold, false)
	k := gocv.GetStructuringElement(gocv.MorphRect, image.Pt(3, 3))
	var windows []*gocv.Window
	if debug {
		windows = []*gocv.Window{gocv.NewWindow("MOG: Bounding boxes"), gocv.NewWindow("MOG: Motion")}
	}
	hold := make([][]byte, hf-1)
	return &MOGFilter{dst, area, &bs, k, debug, windows, hold, hf}
}

// Implements io.Closer.
// Close frees resources used by gocv, because it has to be done manually, due to
// it using c-go.
func (m *MOGFilter) Close() error {
	m.bs.Close()
	m.knl.Close()
	for _, window := range m.windows {
		window.Close()
	}
	return nil
}

// Implements io.Writer.
// Write applies the motion filter to the video stream. Only frames with motion
// are written to the destination encoder, frames without are discarded.
func (m *MOGFilter) Write(f []byte) (int, error) {
	if hfCount < (m.hf - 1) {
		m.hold[hfCount] = f
		hfCount++
		return 0, nil
	} else {
		img, err := gocv.IMDecode(f, gocv.IMReadColor)
		if err != nil {
			return 0, fmt.Errorf("image can't be decoded: %w", err)
		}
		defer img.Close()

		imgDelta := gocv.NewMat()
		defer imgDelta.Close()

		// Seperate foreground and background.
		m.bs.Apply(img, &imgDelta)

		// Threshold imgDelta.
		gocv.Threshold(imgDelta, &imgDelta, 25, 255, gocv.ThresholdBinary)

		// Remove noise.
		gocv.Erode(imgDelta, &imgDelta, m.knl)
		gocv.Dilate(imgDelta, &imgDelta, m.knl)

		// Fill small holes.
		gocv.Dilate(imgDelta, &imgDelta, m.knl)
		gocv.Erode(imgDelta, &imgDelta, m.knl)

		// Find contours and reject ones with a small area.
		var contours [][]image.Point
		allContours := gocv.FindContours(imgDelta, gocv.RetrievalExternal, gocv.ChainApproxSimple)
		for _, c := range allContours {
			if gocv.ContourArea(c) > m.area {
				contours = append(contours, c)
			}
		}

		// Draw debug information.
		if m.debug {
			for _, c := range contours {
				rect := gocv.BoundingRect(c)
				gocv.Rectangle(&img, rect, color.RGBA{0, 0, 255, 0}, 1)
			}

			if len(contours) > 0 {
				gocv.PutText(&img, "Motion", image.Pt(32, 32), gocv.FontHersheyPlain, 2.0, color.RGBA{255, 0, 0, 0}, 2)
			}

			m.windows[0].IMShow(img)
			m.windows[1].IMShow(imgDelta)
			m.windows[0].WaitKey(1)
		}

		// Don't write to destination if there is no motion.
		if len(contours) == 0 {
			return 0, nil
		}

		// Write to destination, past 4 frames then current frame.
		for i, h := range m.hold {
			_, err := m.dst.Write(h)
			m.hold[i] = nil
			if err != nil {
				return 0, err
			}
		}

		hfCount = 0
		return m.dst.Write(f)
	}
}