av/filter/knn.go

// +build !circleci

/*
DESCRIPTION
  A filter that detects motion and discards frames without motion. The
  filter uses a K-Nearest Neighbours (KNN) to determine what is
  background and what is foreground.

AUTHORS
  Ella Pietraroia <ella@ausocean.org>

LICENSE
  KNN.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"
	"io"

	"bitbucket.org/ausocean/av/revid/config"
	"gocv.io/x/gocv"
)

const (
	defaultKNNMinArea     = 25.0
	defaultKNNThreshold   = 300
	defaultKNNHistory     = 300
	defaultKNNKernel      = 4
	defaultKNNDownscaling = 2
	defaultKNNInterval    = 1
)

// KNN is a filter that provides basic motion detection. KNN is short for
// K-Nearest Neighbours method.
type KNN struct {
	debugging debugWindows
	dst       io.WriteCloser                // Destination to which motion containing frames go.
	area      float64                       // The minimum area that a contour can be found in.
	bs        *gocv.BackgroundSubtractorKNN // Uses the KNN algorithm to find the difference between the current and background frame.
	knl       gocv.Mat                      // Matrix that is used for calculations.
	hold      [][]byte                      // Will hold all frames up to hf (so only every hf frame is motion detected).
	hf        int                           // The number of frames to be held.
	hfCount   int                           // Counter for the hold array.
	scale     float64                       // The factor that frames will be downscaled by for motion detection.
}

// NewKNN returns a pointer to a new KNN filter struct.
func NewKNN(dst io.WriteCloser, c config.Config) *KNN {
	// Validate parameters.
	if c.MotionMinArea <= 0 {
		c.LogInvalidField("MotionMinArea", defaultKNNMinArea)
		c.MotionMinArea = defaultKNNMinArea
	}
	if c.MotionThreshold <= 0 {
		c.LogInvalidField("MotionThreshold", defaultKNNThreshold)
		c.MotionThreshold = defaultKNNThreshold
	}
	if c.MotionHistory == 0 {
		c.LogInvalidField("MotionHistory", defaultKNNHistory)
		c.MotionHistory = defaultKNNHistory
	}
	if c.MotionDownscaling <= 0 {
		c.LogInvalidField("MotionDownscaling", defaultKNNDownscaling)
		c.MotionDownscaling = defaultKNNDownscaling
	}
	if c.MotionInterval <= 0 {
		c.LogInvalidField("MotionInterval", defaultKNNInterval)
		c.MotionInterval = defaultKNNInterval
	}
	if c.MotionKernel <= 0 {
		c.LogInvalidField("MotionKernel", defaultKNNKernel)
		c.MotionKernel = defaultKNNKernel
	}

	bs := gocv.NewBackgroundSubtractorKNNWithParams(int(c.MotionHistory), c.MotionThreshold, false)
	k := gocv.GetStructuringElement(gocv.MorphRect, image.Pt(int(c.MotionKernel), int(c.MotionKernel)))

	return &KNN{
		dst:       dst,
		area:      c.MotionMinArea,
		bs:        &bs,
		knl:       k,
		hold:      make([][]byte, c.MotionInterval-1),
		hf:        c.MotionInterval,
		scale:     1 / float64(c.MotionDownscaling),
		debugging: newWindows("KNN"),
	}
}

// Implements io.Closer.
// Close frees resources used by gocv, because it has to be done manually, due to
// it using c-go.
func (m *KNN) Close() error {
	m.bs.Close()
	m.knl.Close()
	m.debugging.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 *KNN) Write(f []byte) (int, error) {
	if m.hfCount < (m.hf - 1) {
		m.hold[m.hfCount] = f
		m.hfCount++
		return len(f), nil
	}
	m.hfCount = 0

	img, err := gocv.IMDecode(f, gocv.IMReadColor)
	if err != nil {
		return 0, fmt.Errorf("can't decode image: %w", err)
	}
	defer img.Close()

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

	// Downsize image to speed up calculations.
	gocv.Resize(img, &img, image.Point{}, m.scale, m.scale, gocv.InterpolationNearestNeighbor)

	// 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.
	m.debugging.show(img, imgDelta, len(contours) > 0, &contours)

	// Don't write to destination if there is no motion.
	if len(contours) == 0 {
		return len(f), 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 len(f), fmt.Errorf("could not write previous frames: %w", err)
		}
	}
	return m.dst.Write(f)
}