av/cmd/rv/probe.go

//go:build !nocv
// +build !nocv

/*
DESCRIPTION
  Provides the methods for the turbidity probe using GoCV. Turbidity probe
  will collect the most recent frames in a buffer and write the latest sharpness
  and contrast scores to the probe.

AUTHORS
  Russell Stanley <russell@ausocean.org>

LICENSE
  Copyright (C) 2021-2022 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 main

import (
	"bytes"
	"errors"
	"os"
	"time"

	"bitbucket.org/ausocean/av/turbidity"
	"bitbucket.org/ausocean/utils/logger"
	"gocv.io/x/gocv"
	"gonum.org/v1/gonum/stat"
)

var errNotEnoughBytes = errors.New("not enough bytes to read")

// Misc constants.
const (
	maxImages        = 1 // Max number of images read when evaluating turbidity.
	turbidityTimeout = 50 * time.Second
)

// Turbidity sensor constants.
const (
	k1, k2     = 4, 4 // Block size, must be divisible by the size template with no remainder.
	filterSize = 3    // Sobel filter size.
	scale      = 1.0  // Amount of scale applied to sobel filter values.
	alpha      = 1.0  // Paramater for contrast equation.
)

type turbidityProbe struct {
	sharpness, contrast float64
	delay               time.Duration
	ticker              time.Ticker
	ts                  *turbidity.TurbiditySensor
	log                 logger.Logger
	buffer              *bytes.Buffer
}

type frameScanner struct {
	off int
	buf []byte
}

// NewTurbidityProbe returns a new turbidity probe.
func NewTurbidityProbe(log logger.Logger, delay time.Duration) (*turbidityProbe, error) {
	tp := new(turbidityProbe)
	tp.log = log
	tp.delay = delay
	tp.ticker = *time.NewTicker(delay)
	tp.buffer = bytes.NewBuffer(*new([]byte))

	// Create the turbidity sensor.
	standard := gocv.IMRead("../../turbidity/images/default.jpg", gocv.IMReadGrayScale)
	template := gocv.IMRead("../../turbidity/images/template.jpg", gocv.IMReadGrayScale)
	ts, err := turbidity.NewTurbiditySensor(template, standard, k1, k2, filterSize, scale, alpha, log)
	if err != nil {
		log.Error("failed create turbidity sensor", "error", err.Error())
	}
	tp.ts = ts
	return tp, nil
}

// Write, reads input h264 frames in the form of a byte stream and writes the the sharpness and contrast
// scores of a video to the the turbidity probe.
func (tp *turbidityProbe) Write(p []byte) (int, error) {
	if len(tp.buffer.Bytes()) == 0 {
		video, err := trim(p)
		if err != nil {
			tp.log.Log(logger.Warning, "trim error", "error", err.Error())
			return 0, nil
		}
		tp.buffer.Write(video)
		tp.log.Log(logger.Debug, "video trimmed, write keyframe complete", "size(bytes)", len(tp.buffer.Bytes()))
	} else if len(tp.buffer.Bytes()) < 100000 {
		tp.buffer.Write(p)
		tp.log.Log(logger.Debug, "write to video buffer complete", "size(bytes)", len(tp.buffer.Bytes()))
	}

	select {
	case <-tp.ticker.C:
		tp.buffer.Reset()
		/*
			tp.log.Log(logger.Debug, "beginning turbidity calculation")
			startTime := time.Now()
			tp.turbidityCalculation()
			tp.log.Log(logger.Debug, "finished turbidity calculation", "total duration (sec)", time.Since(startTime).Seconds())
		*/
	default:
	}
	return len(p), nil
}

func (tp *turbidityProbe) Close() error {
	return nil
}

func (tp *turbidityProbe) turbidityCalculation() {
	var imgs []gocv.Mat
	img := gocv.NewMat()
	// Write byte array to a temp file.
	file, err := os.CreateTemp("temp", "video*.h264")
	if err != nil {
		tp.log.Error("failed to create temp file", "error", err.Error())
		// TODO: Error handling.
		return
	}
	startTime := time.Now()
	_, err = file.Write(tp.buffer.Bytes()[:])
	if err != nil {
		tp.log.Error("failed to write to temporary file", "error", err.Error())
		// TODO: Error handling.
		return
	}
	tp.log.Log(logger.Debug, "writing to temp file", "total duration (sec)", time.Since(startTime).Seconds())
	tp.log.Log(logger.Debug, "buffer", "size(bytes)", len(tp.buffer.Bytes()))
	tp.buffer.Reset()

	startTime = time.Now()
	// Read the file and store each frame.
	vc, err := gocv.VideoCaptureFile(file.Name())
	if err != nil {
		tp.log.Error("failed to open video file", "error", err.Error())
		// TODO: Error handling.
		return
	}
	tp.log.Log(logger.Debug, ".h264 decoded", "total duration (sec)", time.Since(startTime).Seconds())
	startTime = time.Now()
	for vc.Read(&img) && len(imgs) < maxImages {
		imgs = append(imgs, img.Clone())
	}
	if len(imgs) <= 0 {
		tp.log.Log(logger.Warning, "no frames found")
		return
	}
	tp.log.Log(logger.Debug, "read time", "total duration (sec)", time.Since(startTime).Seconds())
	// Process video data to get saturation and contrast scores.
	startTime = time.Now()
	res, err := tp.ts.Evaluate(imgs)
	if err != nil {
		tp.log.Error("evaluate failed", "error", err.Error())
		// TODO: Error handling.
	} else {
		tp.contrast = stat.Mean(res.Contrast, nil)
		tp.sharpness = stat.Mean(res.Sharpness, nil)
	}
	tp.log.Log(logger.Debug, "evaluate complete", "total duration (sec)", time.Since(startTime).Seconds())

	// Cleanup
	err = os.Remove(file.Name())
	if err != nil {
		tp.log.Error("could not remove file", "error", err.Error())
	}
	err = vc.Close()
	if err != nil {
		tp.log.Error("could not close video capture", "error", err.Error())
	}
}

func trim(n []byte) ([]byte, error) {
	sc := frameScanner{buf: n}
	for {
		b, ok := sc.readByte()
		if !ok {
			return nil, errNotEnoughBytes
		}
		for i := 1; b == 0x00 && i != 4; i++ {
			b, ok = sc.readByte()
			if !ok {
				return nil, errNotEnoughBytes
			}
			if b != 0x01 || (i != 2 && i != 3) {
				continue
			}

			b, ok = sc.readByte()
			if !ok {
				return nil, errNotEnoughBytes
			}
			nalType := int(b & 0x1f)
			if nalType == 7 {
				sc.off = sc.off - 4
				return sc.buf[sc.off:], nil
			}
		}
	}
}

func (s *frameScanner) readByte() (b byte, ok bool) {
	if s.off >= len(s.buf) {
		return 0, false
	}
	b = s.buf[s.off]
	s.off++
	return b, true
}