av/filter/basic.go

/*
DESCRIPTION
  A filter that detects motion and discards frames without motion. The
  filter uses a difference method looking at each individual pixel to
  determine what is background and what is foreground.

AUTHORS
  Ella Pietraroia <ella@ausocean.org>

LICENSE
  basic.go is Copyright (C) 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

import (
	"bytes"
	"fmt"
	"image"
	"image/color"
	"image/jpeg"
	"io"
	"os"
	"strconv"
	"sync"

	"golang.org/x/image/font"
	"golang.org/x/image/font/basicfont"
	"golang.org/x/image/math/fixed"
)

const debugfile = "motion.mjpeg"

type pixel struct{ r, g, b uint32 }

// Basic is a filter that provides basic motion detection via a difference
// method.
type Basic struct {
	dst    io.WriteCloser
	img    image.Image
	bg     [][]pixel
	bwImg  *image.RGBA
	thresh int
	pix    int
	w      int
	h      int
	file   io.WriteCloser
	motion int
	debug  bool
}

// NewBasic returns a pointer to a new Basic filter struct.
func NewBasic(dst io.WriteCloser, debug bool, t, p int) *Basic {
	bwImg := image.NewRGBA(image.Rect(0, 0, 0, 0))
	var file io.WriteCloser
	var err error
	file = nil
	if debug {
		file, err = os.Create(debugfile)
		if err != nil {
			panic(fmt.Sprintf("could not create debug file: %v", err))
		}
	}
	return &Basic{dst, nil, nil, bwImg, t, p, 0, 0, file, 0, debug}
}

// Implements io.Closer.
func (bf *Basic) Close() error {
	if bf.debug {
		bf.file.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 (bf *Basic) Write(f []byte) (int, error) {
	// Decode MJPEG.
	var err error
	bf.img, err = jpeg.Decode(bytes.NewReader(f))
	if err != nil {
		return 0, fmt.Errorf("image can't be decoded: %w", err)
	}

	// First frame must be set as the first background image
	if bf.bg == nil {
		bounds := bf.img.Bounds()
		bf.w = bounds.Max.X
		bf.h = bounds.Max.Y

		bf.bwImg = image.NewRGBA(image.Rect(0, 0, bf.w, bf.h))

		bf.bg = make([][]pixel, bf.h)
		for j, _ := range bf.bg {
			bf.bg[j] = make([]pixel, bf.w)
			for i, _ := range bf.bg[j] {
				p := bf.img.At(i, j)
				r, g, b, _ := p.RGBA()
				bf.bg[j][i].r = r
				bf.bg[j][i].b = b
				bf.bg[j][i].g = g
			}
		}
		return len(f), nil
	}

	// Use 4x goroutines to each process one row of pixels.
	var j int
	j = 0
	var wg sync.WaitGroup

	for j < bf.h {
		wg.Add(4)
		go bf.process(j, &wg)
		go bf.process(j+1, &wg)
		go bf.process(j+2, &wg)
		go bf.process(j+3, &wg)
		j = j + 4
		wg.Wait()
	}

	// Will save a video of where motion is detected in motion.mjpeg (in the current folder).
	if bf.debug {
		err := bf.saveFrame()
		if err != nil {
			return len(f), fmt.Errorf("image can't be encoded for debug video: %w", err)
		}
	}

	// If there are not enough motion pixels then discard the frame.
	if bf.motion < pixels {
		return len(f), nil
	}

	// Write all motion frames.
	return bf.dst.Write(f)
}

// Go routine for one row of the image to be processed
func (bf *Basic) process(j int, wg *sync.WaitGroup) {
	for i, _ := range bf.bg[j] {
		n := bf.img.At(i, j)
		r, b, g, _ := n.RGBA()

		// Compare the difference of the RGB values of each pixel to the background image.
		diffR := absDiff(r, bf.bg[j][i].r)
		diffG := absDiff(g, bf.bg[j][i].g)
		diffB := absDiff(g, bf.bg[j][i].b)

		diff := diffR + diffG + diffB

		if diff > threshold {
			bf.motion++
		}

		if bf.debug {
			if diff > threshold {
				bf.bwImg.SetRGBA(i, j, color.RGBA{0xff, 0xff, 0xff, 0xff})
			} else {
				bf.bwImg.SetRGBA(i, j, color.RGBA{0x00, 0x00, 0x00, 0xff})
			}
		}

		// Update backgound image.
		bf.bg[j][i].r = r
		bf.bg[j][i].b = b
		bf.bg[j][i].g = g
	}
	wg.Done()
}

func (bf *Basic) saveFrame() error {
	col := color.RGBA{200, 100, 0, 255}
	d := &font.Drawer{
		Dst:  bf.bwImg,
		Src:  image.NewUniform(col),
		Face: basicfont.Face7x13,
		Dot:  fixed.P(20, 30),
	}
	var s string
	if bf.motion > bf.pix {
		s = strconv.Itoa(bf.motion) + " Motion"
	} else {
		s = strconv.Itoa(bf.motion)
	}
	d.DrawString(s)
	return jpeg.Encode(bf.file, bf.bwImg, nil)
}

func absDiff(a, b uint32) int {
	c := int(a) - int(b)
	if c < 0 {
		return -c
	} else {
		return c
	}

}