av/filter/filter.go

/*
NAME
  filter.go

AUTHORS
  Ella Pietraroia <ella@ausocean.org>
  Scott Barnard <scott@ausocean.org>

LICENSE
  filter.go is Copyright (C) 2017-2020 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 provides filters that modify video streams.
package filter

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

	"gocv.io/x/gocv"
)

// Filter is an interface shared by all filters. It is the same as an io.WriteCloser
// interface, however more functions may be added as needed.
type Filter interface {
	io.WriteCloser
}

// MogFilter is a filter that provides basic motion detection.
type MogFilter struct {
	dst     io.WriteCloser
	minArea float64
	mog2    *gocv.BackgroundSubtractorMOG2
	kernel  gocv.Mat
	debug   bool
	windows []*gocv.Window
}

// NewMogFilter returns a pointer to a new MogFilter
func NewMogFilter(dst io.WriteCloser, minArea, threshold float64, history, kernelSize int, debug bool) *MogFilter {
	mog2 := gocv.NewBackgroundSubtractorMOG2WithParams(history, threshold, false)
	kernel := gocv.GetStructuringElement(gocv.MorphRect, image.Pt(kernelSize, kernelSize))
	var windows []*gocv.Window
	if debug {
		windows = append(windows, gocv.NewWindow("Debug: Bounding boxes"), gocv.NewWindow("Debug: Motion"))
	}
	return &MogFilter{dst, minArea, &mog2, kernel, debug, windows}
}

// 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.mog2.Close()
	m.kernel.Close()
	for _, window := range m.windows {
		window.Close()
	}
	return m.dst.Close()
}

// 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) {
	img, _ := gocv.IMDecode(f, gocv.IMReadColor)
	defer img.Close()

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

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

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

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

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

	// 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.minArea {
			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 -1, nil
	}

	// Write to destination.
	return m.dst.Write(f)
}