av/cmd/rv/main.go

/*
DESCRIPTION
  rv is a netsender client using the revid package to perform media collection
  and forwarding whose behaviour is controllable via the cloud interfaces
  netreceiver and vidgrind.

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>
  Alan Noble <alan@ausocean.org>
  Dan Kortschak <dan@ausocean.org>
  Jack Richardson <jack@ausocean.org>
  Trek Hopton <trek@ausocean.org>
  Scott Barnard <scott@ausocean.org>
  Russell Stanley <russell@ausocean.org>

LICENSE
  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
  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.

USAGE
  There must firstly be a netsender configuration file under /etc/netsender.conf.
  Example:

  ma 00:00:00:00:00:01
  dk 0
  wi
  ip V0, T0
  op
  mp 60
  ap 0
  tg
  hw
  sh vidgrind.appspot.com

  Revid configuration is controlled by valid variables given values on netreceiver
  or vidgrind interface. See revid/config for valid variables.

  To run rv simply build and call:
  ./rv
*/

// Package rv is a netsender client for revid.
package main

import (
	"bytes"
	"io"
	"os"
	"runtime/pprof"
	"strconv"
	"time"

	"gocv.io/x/gocv"
	"gonum.org/v1/gonum/stat"
	"gopkg.in/natefinch/lumberjack.v2"

	"bitbucket.org/ausocean/av/container/mts"
	"bitbucket.org/ausocean/av/container/mts/meta"
	"bitbucket.org/ausocean/av/revid"
	"bitbucket.org/ausocean/av/revid/config"
	"bitbucket.org/ausocean/av/turbidity"
	"bitbucket.org/ausocean/iot/pi/netlogger"
	"bitbucket.org/ausocean/iot/pi/netsender"
	"bitbucket.org/ausocean/utils/logger"
)

// Copyright information prefixed to all metadata.
const (
	metaPreambleKey  = "copyright"
	metaPreambleData = "ausocean.org/license/content2019"
)

// Logging configuration.
const (
	logPath      = "/var/log/netsender/netsender.log"
	logMaxSize   = 500 // MB
	logMaxBackup = 10
	logMaxAge    = 28 // days
	logVerbosity = logger.Info
	logSuppress  = true
)

// Revid modes.
const (
	modeNormal = "Normal"
	modePaused = "Paused"
	modeBurst  = "Burst"
	modeLoop   = "Loop"
)

// Misc constants.
const (
	netSendRetryTime = 5 * time.Second
	defaultSleepTime = 60 // Seconds
	profilePath      = "rv.prof"
	pkg              = "rv: "
	runPreDelay      = 20 * time.Second
	maxImages        = 1 // Max number of images read when evaluating turbidity.
)

// Software define pin values.
// See https://netreceiver.appspot.com/help, External Pin Assignments.
const (
	bitratePin   = "X36"
	sharpnessPin = "X38"
	contrastPin  = "X39"
)

// 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.
)

// This is set to true if the 'profile' build tag is provided on build.
var canProfile = false

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

// 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/template.jpg", gocv.IMReadGrayScale)
	template := gocv.IMRead("../../turbidity/images/icon.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) {
	tp.buffer.Write(p)
	go func() {
		select {
		case <-tp.ticker.C:
			tp.turbidityCalculation(p)
		default:
			return
		}
	}()
	return len(p), nil
}

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

func (tp *turbidityProbe) turbidityCalculation(p []byte) {
	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
	}
	defer os.Remove(file.Name())
	_, 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.buffer.Reset()
	// 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
	}
	for vc.Read(&img) && len(imgs) < maxImages {
		imgs = append(imgs, img.Clone())
	}
	if len(imgs) <= 0 {
		tp.log.Log(logger.Warning, "no frames found", "error", err.Error())
		return
	}
	tp.log.Log(logger.Debug, "found frames", "frames", len(imgs))

	// 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.log.Log(logger.Debug, "finished evaluation", "total duration (sec)", time.Since(startTime).Seconds())
		tp.contrast = stat.Mean(res.Contrast, nil)
		tp.sharpness = stat.Mean(res.Sharpness, nil)
	}
	return
}

func main() {
	mts.Meta = meta.NewWith([][2]string{{metaPreambleKey, metaPreambleData}})

	// Create lumberjack logger to handle logging to file.
	fileLog := &lumberjack.Logger{
		Filename:   logPath,
		MaxSize:    logMaxSize,
		MaxBackups: logMaxBackup,
		MaxAge:     logMaxAge,
	}

	// Create netlogger to handle logging to cloud.
	netLog := netlogger.New()

	// Create logger that we call methods on to log, which in turn writes to the
	// lumberjack and netloggers.
	log := logger.New(logVerbosity, io.MultiWriter(fileLog, netLog), logSuppress)

	// If rv has been built with the profile tag, then we'll start a CPU profile.
	if canProfile {
		profile(log)
		defer pprof.StopCPUProfile()
		log.Log(logger.Info, "profiling started")
	}

	var (
		rv *revid.Revid
		p  *turbidityProbe
	)
	p, err := NewTurbidityProbe(*log, 60*time.Second)

	log.Log(logger.Debug, "initialising netsender client")
	ns, err := netsender.New(log, nil, readPin(p, rv, log), nil, createVarMap())
	if err != nil {
		log.Log(logger.Fatal, pkg+"could not initialise netsender client: "+err.Error())
	}

	log.Log(logger.Debug, "initialising revid")
	rv, err = revid.New(config.Config{Logger: log}, ns)
	if err != nil {
		log.Log(logger.Fatal, pkg+"could not initialise revid", "error", err.Error())
	}

	err = rv.SetProbe(p)
	if err != nil {
		log.Log(logger.Error, pkg+"could not set probe", "error", err.Error())
	}

	// NB: Problems were encountered with communicating with RTSP inputs. When trying to
	// connect it would fail due to timeout; as if things had not been set up quickly
	// enough before revid tried to do things. This delay fixes this, but there is probably
	// a better way to solve this problem.
	time.Sleep(runPreDelay)

	log.Log(logger.Debug, "beginning main loop")
	run(rv, ns, log, netLog)
}

// run starts the main loop. This will run netsender on every pass of the loop
// (sleeping inbetween), check vars, and if changed, update revid as appropriate.
func run(rv *revid.Revid, ns *netsender.Sender, l *logger.Logger, nl *netlogger.Logger) {
	var vs int
	for {
		l.Log(logger.Debug, "running netsender")
		err := ns.Run()
		if err != nil {
			l.Log(logger.Warning, pkg+"Run Failed. Retrying...", "error", err.Error())
			time.Sleep(netSendRetryTime)
			continue
		}

		l.Log(logger.Debug, "sending logs")
		err = nl.Send(ns)
		if err != nil {
			l.Log(logger.Warning, pkg+"Logs could not be sent", "error", err.Error())
		}

		l.Log(logger.Debug, "checking varsum")
		newVs := ns.VarSum()
		if vs == newVs {
			sleep(ns, l)
			continue
		}
		vs = newVs
		l.Log(logger.Info, "varsum changed", "vs", vs)

		l.Log(logger.Debug, "getting new vars")
		vars, err := ns.Vars()
		if err != nil {
			l.Log(logger.Error, pkg+"netSender failed to get vars", "error", err.Error())
			time.Sleep(netSendRetryTime)
			continue
		}
		l.Log(logger.Debug, "got new vars", "vars", vars)

		// Configure revid based on the vars.
		l.Log(logger.Debug, "updating revid's configuration")
		err = rv.Update(vars)
		if err != nil {
			l.Log(logger.Warning, pkg+"couldn't update revid", "error", err.Error())
			sleep(ns, l)
			continue
		}
		l.Log(logger.Info, "revid successfully reconfigured")

		l.Log(logger.Debug, "checking mode")
		switch ns.Mode() {
		case modePaused:
			l.Log(logger.Debug, "mode is Paused, stopping revid")
			rv.Stop()
		case modeNormal, modeLoop:
			l.Log(logger.Debug, "mode is Normal or Loop, starting revid")
			err = rv.Start()
			if err != nil {
				l.Log(logger.Error, pkg+"could not start revid", "error", err.Error())
				ns.SetMode(modePaused, &vs)
				sleep(ns, l)
				continue
			}
		case modeBurst:
			l.Log(logger.Debug, "mode is Burst, bursting revid")
			err = rv.Burst()
			if err != nil {
				l.Log(logger.Warning, pkg+"could not start burst", "error", err.Error())
				ns.SetMode(modePaused, &vs)
				sleep(ns, l)
				continue
			}
			ns.SetMode(modePaused, &vs)
		}
		l.Log(logger.Info, "revid updated with new mode")

		sleep(ns, l)
	}
}

func createVarMap() map[string]string {
	m := make(map[string]string)
	for _, v := range config.Variables {
		m[v.Name] = v.Type
	}
	return m
}

// profile opens a file to hold CPU profiling metrics and then starts the
// CPU profiler.
func profile(l *logger.Logger) {
	f, err := os.Create(profilePath)
	if err != nil {
		l.Log(logger.Fatal, pkg+"could not create CPU profile", "error", err.Error())
	}
	if err := pprof.StartCPUProfile(f); err != nil {
		l.Log(logger.Fatal, pkg+"could not start CPU profile", "error", err.Error())
	}
}

// sleep uses a delay to halt the program based on the monitoring period
// netsender parameter (mp) defined in the netsender.conf config.
func sleep(ns *netsender.Sender, l *logger.Logger) {
	l.Log(logger.Debug, "sleeping")
	t, err := strconv.Atoi(ns.Param("mp"))
	if err != nil {
		l.Log(logger.Error, pkg+"could not get sleep time, using default", "error", err)
		t = defaultSleepTime
	}
	time.Sleep(time.Duration(t) * time.Second)
	l.Log(logger.Debug, "finished sleeping")
}

// readPin provides a callback function of consistent signature for use by
// netsender to retrieve software defined pin values e.g. revid bitrate.
func readPin(p *turbidityProbe, rv *revid.Revid, l *logger.Logger) func(pin *netsender.Pin) error {
	return func(pin *netsender.Pin) error {
		switch {
		case pin.Name == bitratePin:
			pin.Value = -1
			if rv != nil {
				pin.Value = rv.Bitrate()
			}
		case pin.Name == sharpnessPin:
			pin.Value = -1
			if p != nil {
				l.Debug("setting sharpness value", "sharpness", p.sharpness*1000)
				pin.Value = int(p.sharpness * 1000)
			}
		case pin.Name == contrastPin:
			pin.Value = -1
			if p != nil {
				l.Debug("setting contrast pin", "contrast", p.contrast)
				pin.Value = int(p.contrast * 100)
			}
		default:
			pin.Value = -1
		}
		return nil
	}
}