av/cmd/revid-cli/main.go

/*
NAME
  revid-cli - command line interface for Revid.

DESCRIPTION
  See Readme.md

AUTHORS
  Saxon A. Nelson-Milton <saxon@ausocean.org>
  Jack Richardson <jack@ausocean.org>

LICENSE
  revid-cli is Copyright (C) 2017-2018 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.
*/

package main

import (
	"flag"
	"os"
	"runtime/pprof"
	"strconv"
	"strings"
	"time"

	"bitbucket.org/ausocean/av/container/mts"
	"bitbucket.org/ausocean/av/container/mts/meta"
	"bitbucket.org/ausocean/av/revid"
	"bitbucket.org/ausocean/iot/pi/netsender"
	"bitbucket.org/ausocean/iot/pi/sds"
	"bitbucket.org/ausocean/iot/pi/smartlogger"
	"bitbucket.org/ausocean/utils/logger"
)

// Revid modes
const (
	normal = "Normal"
	paused = "Paused"
	burst  = "Burst"
)

// Other misc consts
const (
	netSendRetryTime    = 5 * time.Second
	defaultRunDuration  = 24 * time.Hour
	revidStopTime       = 5 * time.Second
	defaultLogPath      = "/var/log/netsender"
	pkg                 = "revid-cli:"
	defaultLogVerbosity = logger.Info
	defaultSleepTime    = 60 // Seconds
)

// canProfile is set to false with revid-cli is built with "-tags profile".
var canProfile = true

// The logger that will be used throughout.
var log *logger.Logger

const (
	metaPreambleKey  = "copyright"
	metaPreambleData = "ausocean.org/license/content2019"
)

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

	useNetsender := flag.Bool("NetSender", false, "Are we checking vars through netsender?")
	runDurationPtr := flag.Duration("runDuration", defaultRunDuration, "How long do you want revid to run for?")

	cfg := handleFlags()
	if !*useNetsender {
		rv, err := revid.New(cfg, nil)
		if err != nil {
			cfg.Logger.Log(logger.Fatal, pkg+"failed to initialiase revid", "error", err.Error())
		}
		if err = rv.Start(); err != nil {
			cfg.Logger.Log(logger.Fatal, pkg+"failed to start revid", "error", err.Error())
		}
		time.Sleep(*runDurationPtr)
		rv.Stop()
		return
	}

	run(cfg)
}

// handleFlags parses command line flags and returns a revid configuration
// based on them.
func handleFlags() revid.Config {
	var cfg revid.Config

	var (
		cpuprofile = flag.String("cpuprofile", "", "write cpu profile to `file`")

		inputPtr              = flag.String("Input", "", "The input type: Raspivid, File, v4l, Audio")
		inputCodecPtr         = flag.String("InputCodec", "", "The codec of the input: H264, Mjpeg, PCM, ADPCM")
		rtmpMethodPtr         = flag.String("RtmpMethod", "", "The method used to send over rtmp: Ffmpeg, Librtmp")
		quantizePtr           = flag.Bool("Quantize", false, "Quantize input (non-variable bitrate)")
		verbosityPtr          = flag.String("Verbosity", "Info", "Verbosity: Debug, Info, Warning, Error, Fatal")
		rtpAddrPtr            = flag.String("RtpAddr", "", "Rtp destination address: <IP>:<port> (port is generally 6970-6999)")
		logPathPtr            = flag.String("LogPath", defaultLogPath, "The log path")
		configFilePtr         = flag.String("ConfigFile", "", "NetSender config file")
		rtmpUrlPtr            = flag.String("RtmpUrl", "", "Url of rtmp endpoint")
		outputPathPtr         = flag.String("OutputPath", "", "The directory of the output file")
		inputFilePtr          = flag.String("InputPath", "", "The directory of the input file")
		httpAddressPtr        = flag.String("HttpAddress", "", "Destination address of http posts")
		sendRetryPtr          = flag.Bool("retry", false, "Specify whether a failed send should be retried.")
		verticalFlipPtr       = flag.Bool("VerticalFlip", false, "Flip video vertically: Yes, No")
		horizontalFlipPtr     = flag.Bool("HorizontalFlip", false, "Flip video horizontally: Yes, No")
		framesPerClipPtr      = flag.Uint("FramesPerClip", 0, "Number of frames per clip sent")
		bitratePtr            = flag.Uint("Bitrate", 0, "Bitrate of recorded video")
		heightPtr             = flag.Uint("Height", 0, "Height in pixels")
		widthPtr              = flag.Uint("Width", 0, "Width in pixels")
		frameRatePtr          = flag.Uint("FrameRate", 0, "Frame rate of captured video")
		quantizationPtr       = flag.Uint("Quantization", 0, "Desired quantization value: 0-40")
		intraRefreshPeriodPtr = flag.Uint("IntraRefreshPeriod", 0, "The IntraRefreshPeriod i.e. how many keyframes we send")
		rotationPtr           = flag.Uint("Rotation", 0, "Rotate video output. (0-359 degrees)")
		brightnessPtr         = flag.Uint("Brightness", 50, "Set brightness. (0-100) ")
		saturationPtr         = flag.Int("Saturation", 0, "Set Saturation. (100-100)")
		exposurePtr           = flag.String("Exposure", "auto", "Set exposure mode. ("+strings.Join(revid.ExposureModes[:], ",")+")")
		autoWhiteBalancePtr   = flag.String("Awb", "auto", "Set automatic white balance mode. ("+strings.Join(revid.AutoWhiteBalanceModes[:], ",")+")")

		// Audio specific flags.
		sampleRatePtr = flag.Int("SampleRate", 48000, "Sample rate of recorded audio")
		channelsPtr   = flag.Int("Channels", 1, "Record in Mono or Stereo (1 or 2)")
		recPeriodPtr  = flag.Int("recPeriod", 1, "How many seconds to record at a time")
		bitDepthPtr   = flag.Int("bitDepth", 16, "Bit Depth to record audio at.")
	)

	var outputs flagStrings
	flag.Var(&outputs, "Output", "output type: Http, Rtmp, File, Udp, Rtp (may be used more than once)")

	flag.Parse()

	switch *verbosityPtr {
	case "Debug":
		cfg.LogLevel = logger.Debug
	case "Info":
		cfg.LogLevel = logger.Info
	case "Warning":
		cfg.LogLevel = logger.Warning
	case "Error":
		cfg.LogLevel = logger.Error
	case "Fatal":
		cfg.LogLevel = logger.Fatal
	default:
		cfg.LogLevel = defaultLogVerbosity
	}

	log = logger.New(cfg.LogLevel, &smartlogger.New(*logPathPtr).LogRoller)

	cfg.Logger = log

	if *cpuprofile != "" {
		if canProfile {
			f, err := os.Create(*cpuprofile)
			if err != nil {
				log.Log(logger.Fatal, pkg+"could not create CPU profile", "error", err.Error())
			}
			if err := pprof.StartCPUProfile(f); err != nil {
				log.Log(logger.Fatal, pkg+"could not start CPU profile", "error", err.Error())
			}
			defer pprof.StopCPUProfile()
		} else {
			log.Log(logger.Warning, pkg+"ignoring cpuprofile flag - http/pprof built in.")
		}
	}

	switch *inputPtr {
	case "Raspivid":
		cfg.Input = revid.Raspivid
	case "v4l":
		cfg.Input = revid.V4L
	case "File":
		cfg.Input = revid.File
	case "Audio":
		cfg.Input = revid.Audio
	case "":
	default:
		log.Log(logger.Error, pkg+"bad input argument")
	}

	switch *inputCodecPtr {
	case "H264":
		cfg.InputCodec = revid.H264
	case "PCM":
		cfg.InputCodec = revid.PCM
	case "ADPCM":
		cfg.InputCodec = revid.ADPCM
	case "":
	default:
		log.Log(logger.Error, pkg+"bad input codec argument")
	}

	switch *inputPtr {
	case "Audio":
		cfg.WriteRate = uint(*recPeriodPtr)
	default:
		cfg.WriteRate = *frameRatePtr
	}

	if len(outputs) == 0 {
		cfg.Outputs = make([]uint8, 1)
	}

	for _, o := range outputs {
		switch o {
		case "File":
			cfg.Outputs = append(cfg.Outputs, revid.File)
		case "Http":
			cfg.Outputs = append(cfg.Outputs, revid.Http)
		case "Rtmp":
			cfg.Outputs = append(cfg.Outputs, revid.Rtmp)
		case "Rtp":
			cfg.Outputs = append(cfg.Outputs, revid.Rtp)
		case "":
		default:
			log.Log(logger.Error, pkg+"bad output argument", "arg", o)
		}
	}

	switch *rtmpMethodPtr {
	case "Ffmpeg":
		cfg.RtmpMethod = revid.Ffmpeg
	case "LibRtmp":
		cfg.RtmpMethod = revid.LibRtmp
	case "":
	default:
		log.Log(logger.Error, pkg+"bad rtmp method argument")
	}

	if *configFilePtr != "" {
		netsender.ConfigFile = *configFilePtr
	}

	cfg.Quantize = *quantizePtr
	cfg.Rotation = *rotationPtr
	cfg.FlipHorizontal = *horizontalFlipPtr
	cfg.FlipVertical = *verticalFlipPtr
	cfg.FramesPerClip = *framesPerClipPtr
	cfg.RtmpUrl = *rtmpUrlPtr
	cfg.Bitrate = *bitratePtr
	cfg.OutputPath = *outputPathPtr
	cfg.InputPath = *inputFilePtr
	cfg.Height = *heightPtr
	cfg.Width = *widthPtr
	cfg.FrameRate = *frameRatePtr
	cfg.HttpAddress = *httpAddressPtr
	cfg.Quantization = *quantizationPtr
	cfg.IntraRefreshPeriod = *intraRefreshPeriodPtr
	cfg.RtpAddress = *rtpAddrPtr
	cfg.SendRetry = *sendRetryPtr
	cfg.Brightness = *brightnessPtr
	cfg.Saturation = *saturationPtr
	cfg.Exposure = *exposurePtr
	cfg.AutoWhiteBalance = *autoWhiteBalancePtr
	cfg.SampleRate = *sampleRatePtr
	cfg.Channels = *channelsPtr
	cfg.RecPeriod = *recPeriodPtr
	cfg.BitDepth = *bitDepthPtr

	return cfg
}

// initialize then run the main NetSender client
func run(cfg revid.Config) {
	log.Log(logger.Info, pkg+"running in NetSender mode")

	var rv *revid.Revid

	readPin := func(pin *netsender.Pin) error {
		switch {
		case pin.Name == "X23":
			pin.Value = -1
			if rv != nil {
				pin.Value = rv.Bitrate()
			}
		case pin.Name[0] == 'X':
			return sds.ReadSystem(pin)
		default:
			pin.Value = -1
		}
		return nil // Return error only if we want NetSender to generate an error
	}

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

	var vs int
	for {
		err = ns.Run()
		if err != nil {
			log.Log(logger.Error, pkg+"Run Failed. Retrying...", "error", err.Error())
			time.Sleep(netSendRetryTime)
			continue
		}

		// If var sum hasn't changed we continue.
		var vars map[string]string
		newVs := ns.VarSum()
		if vs == newVs {
			goto sleep
		}
		vs = newVs

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

		if rv == nil {
			rv, err = revid.New(cfg, ns)
			if err != nil {
				log.Log(logger.Warning, pkg+"could not initialise revid", "error", err.Error())
				goto sleep
			}
		}

		err = rv.Update(vars)
		if err != nil {
			log.Log(logger.Warning, pkg+"Couldn't update revid", "error", err.Error())
			goto sleep
		}

		switch ns.Mode() {
		case paused:
			rv.Stop()
		case normal:
			err = rv.Start()
			if err != nil {
				log.Log(logger.Warning, pkg+"could not start revid", "error", err.Error())
				ns.SetMode(paused, &vs)
				goto sleep
			}
		case burst:
			log.Log(logger.Info, pkg+"Starting burst...")

			err = rv.Start()
			if err != nil {
				log.Log(logger.Warning, pkg+"could not start burst", "error", err.Error())
				ns.SetMode(paused, &vs)
				goto sleep
			}

			time.Sleep(time.Duration(rv.Config().BurstPeriod) * time.Second)
			log.Log(logger.Info, pkg+"Stopping burst...")
			rv.Stop()
			ns.SetMode(paused, &vs)
		}

	sleep:
		sleepTime, err := strconv.Atoi(ns.Param("mp"))
		if err != nil {
			log.Log(logger.Error, pkg+"could not get sleep time, using default")
			sleepTime = defaultSleepTime
		}
		time.Sleep(time.Duration(sleepTime) * time.Second)
	}
}

// flagStrings implements an appending string set flag.
type flagStrings []string

func (v *flagStrings) String() string {
	if *v != nil {
		return strings.Join(*v, ",")
	}
	return ""
}

func (v *flagStrings) Set(s string) error {
	if s == "" {
		return nil
	}
	for _, e := range *v {
		if e == s {
			return nil
		}
	}
	*v = append(*v, s)
	return nil
}

func (v *flagStrings) Get() interface{} { return *v }