av/filter/knn.go

128 lines
3.5 KiB
Go

/*
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 (
"image"
"image/color"
"io"
"gocv.io/x/gocv"
)
// KNNFilter is a filter that provides basic motion detection. KNN is short for
// K-Nearest Neighbours method.
type KNNFilter struct {
dst io.WriteCloser
area float64
bs *gocv.BackgroundSubtractorKNN
knl gocv.Mat
debug bool
windows []*gocv.Window
}
// NewKNNFilter returns a pointer to a new KNNFilter.
func NewKNNFilter(dst io.WriteCloser, area, threshold float64, history, kernelSize int, debug bool) *KNNFilter {
bs := gocv.NewBackgroundSubtractorKNNWithParams(history, threshold, false)
k := gocv.GetStructuringElement(gocv.MorphRect, image.Pt(kernelSize, kernelSize))
var windows []*gocv.Window
if debug {
windows = []*gocv.Window{gocv.NewWindow("Debug: Bounding boxes"), gocv.NewWindow("Debug: Motion")}
}
return &KNNFilter{dst, area, &bs, k, debug, windows}
}
// Implements io.Closer.
// Close frees resources used by gocv, because it has to be done manually, due to
// it using c-go.
func (m *KNNFilter) 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 *KNNFilter) 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.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 -1, nil
}
// Write to destination.
return m.dst.Write(f)
}