av/device/raspivid/raspivid.go

/*
DESCRIPTION
  raspivid.go provides an implementation of the AVDevice interface for raspivid.

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

LICENSE
  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 raspivid

import (
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os/exec"
	"strings"

	"bitbucket.org/ausocean/av/codec/codecutil"
	"bitbucket.org/ausocean/av/device"
	"bitbucket.org/ausocean/av/revid/config"
	"bitbucket.org/ausocean/utils/logger"
	"bitbucket.org/ausocean/utils/sliceutils"
)

// To indicate package when logging.
const pkg = "raspivid: "

// Raspivid configuration defaults.
const (
	defaultRaspividCodec            = codecutil.CodecH264
	defaultRaspividRotation         = 0
	defaultRaspividWidth            = 1280
	defaultRaspividHeight           = 720
	defaultRaspividBrightness       = 50
	defaultRaspividSaturation       = 0
	defaultRaspividExposure         = "auto"
	defaultRaspividAutoWhiteBalance = "auto"
	defaultRaspividMinFrames        = 100
	defaultRaspividQuantization     = 30
	defaultRaspividBitrate          = 4800
	defaultRaspividFramerate        = 25
)

// Configuration errors.
var (
	errBadCodec            = errors.New("codec bad or unset, defaulting")
	errBadRotation         = errors.New("rotation bad or unset, defaulting")
	errBadWidth            = errors.New("width bad or unset, defaulting")
	errBadHeight           = errors.New("height bad or unset, defaulting")
	errBadFrameRate        = errors.New("framerate bad or unset, defaulting")
	errBadBitrate          = errors.New("bitrate bad or unset, defaulting")
	errBadMinFrames        = errors.New("min frames bad or unset, defaulting")
	errBadSaturation       = errors.New("saturation bad or unset, defaulting")
	errBadBrightness       = errors.New("brightness bad or unset, defaulting")
	errBadExposure         = errors.New("exposure bad or unset, defaulting")
	errBadAutoWhiteBalance = errors.New("auto white balance bad or unset, defaulting")
	errBadQuantization     = errors.New("quantization bad or unset, defaulting")
)

// Possible modes for raspivid --exposure parameter.
var ExposureModes = [...]string{
	"auto",
	"night",
	"nightpreview",
	"backlight",
	"spotlight",
	"sports",
	"snow",
	"beach",
	"verylong",
	"fixedfps",
	"antishake",
	"fireworks",
}

// Possible modes for raspivid --awb parameter.
var AutoWhiteBalanceModes = [...]string{
	"off",
	"auto",
	"sun",
	"cloud",
	"shade",
	"tungsten",
	"fluorescent",
	"incandescent",
	"flash",
	"horizon",
}

// Raspivid is an implementation of AVDevice that provides control over the
// raspivid command to allow reading of data from a Raspberry Pi camera.
type Raspivid struct {
	cfg  config.Config
	cmd  *exec.Cmd
	out  io.ReadCloser
	log  config.Logger
	done chan struct{}
}

// New returns a new Raspivid.
func New(l config.Logger) *Raspivid {
	return &Raspivid{
		log:  l,
		done: make(chan struct{}),
	}
}

// Name returns the name of the device.
func (r *Raspivid) Name() string {
	return "Raspivid"
}

// Set will take a Config struct, check the validity of the relevant fields
// and then performs any configuration necessary. If fields are not valid,
// an error is added to the multiError and a default value is used.
func (r *Raspivid) Set(c config.Config) error {
	var errs device.MultiError
	switch c.InputCodec {
	case codecutil.CodecH264, codecutil.CodecMJPEG:
	default:
		c.InputCodec = defaultRaspividCodec
		errs = append(errs, errBadCodec)
	}

	if c.Rotation > 359 {
		c.Rotation = defaultRaspividRotation
		errs = append(errs, errBadRotation)
	}

	if c.Width == 0 {
		c.Width = defaultRaspividWidth
		errs = append(errs, errBadWidth)
	}

	if c.Height == 0 {
		c.Height = defaultRaspividHeight
		errs = append(errs, errBadHeight)
	}

	if c.FrameRate == 0 {
		c.FrameRate = defaultRaspividFramerate
		errs = append(errs, errBadFrameRate)
	}

	if c.CBR || sliceutils.ContainsUint8(c.Outputs, config.OutputRTMP) {
		c.Quantization = 0
		if c.Bitrate <= 0 {
			errs = append(errs, errBadBitrate)
			c.Bitrate = defaultRaspividBitrate
		}
	} else {
		c.Bitrate = 0
		if c.Quantization < 10 || c.Quantization > 40 {
			errs = append(errs, errBadQuantization)
			c.Quantization = defaultRaspividQuantization
		}
	}

	if c.MinFrames <= 0 {
		errs = append(errs, errBadMinFrames)
		c.MinFrames = defaultRaspividMinFrames
	}

	if c.Brightness <= 0 || c.Brightness > 100 {
		errs = append(errs, errBadBrightness)
		c.Brightness = defaultRaspividBrightness
	}

	if c.Saturation < -100 || c.Saturation > 100 {
		errs = append(errs, errBadSaturation)
		c.Saturation = defaultRaspividSaturation
	}

	if c.Exposure == "" || !sliceutils.ContainsString(ExposureModes[:], c.Exposure) {
		errs = append(errs, errBadExposure)
		c.Exposure = defaultRaspividExposure
	}

	if c.AutoWhiteBalance == "" || !sliceutils.ContainsString(AutoWhiteBalanceModes[:], c.AutoWhiteBalance) {
		errs = append(errs, errBadAutoWhiteBalance)
		c.AutoWhiteBalance = defaultRaspividAutoWhiteBalance
	}

	r.cfg = c
	return errs
}

// Start will prepare the arguments for the raspivid command using the
// configuration set using the Set method then call the raspivid command,
// piping the video output from which the Read method will read from.
func (r *Raspivid) Start() error {
	const disabled = "0"
	args := []string{
		"--output", "-",
		"--nopreview",
		"--timeout", disabled,
		"--width", fmt.Sprint(r.cfg.Width),
		"--height", fmt.Sprint(r.cfg.Height),
		"--bitrate", fmt.Sprint(r.cfg.Bitrate * 1000), // Convert from kbps to bps.
		"--framerate", fmt.Sprint(r.cfg.FrameRate),
		"--rotation", fmt.Sprint(r.cfg.Rotation),
		"--brightness", fmt.Sprint(r.cfg.Brightness),
		"--saturation", fmt.Sprint(r.cfg.Saturation),
		"--exposure", fmt.Sprint(r.cfg.Exposure),
		"--awb", fmt.Sprint(r.cfg.AutoWhiteBalance),
	}

	if r.cfg.HorizontalFlip {
		args = append(args, "--hflip")
	}

	if r.cfg.VerticalFlip {
		args = append(args, "--vflip")
	}
	if r.cfg.HorizontalFlip {
		args = append(args, "--hflip")
	}

	switch r.cfg.InputCodec {
	default:
		return fmt.Errorf("revid: invalid input codec: %v", r.cfg.InputCodec)
	case codecutil.CodecH264:
		args = append(args,
			"--codec", "H264",
			"--inline",
			"--intra", fmt.Sprint(r.cfg.MinFrames),
		)
		if !r.cfg.CBR {
			args = append(args, "-qp", fmt.Sprint(r.cfg.Quantization))
		}
	case codecutil.CodecMJPEG:
		args = append(args, "--codec", "MJPEG")
	}
	r.cfg.Logger.Log(logger.Info, pkg+"raspivid args", "raspividArgs", strings.Join(args, " "))
	r.cmd = exec.Command("raspivid", args...)

	var err error
	r.out, err = r.cmd.StdoutPipe()
	if err != nil {
		return fmt.Errorf("could not pipe command output: %w", err)
	}

	stderr, err := r.cmd.StderrPipe()
	if err != nil {
		return fmt.Errorf("could not pipe command error: %w", err)
	}

	go func() {
		for {
			select {
			case <-r.done:
				r.cfg.Logger.Log(logger.Info, "raspivid.Stop() called, finished checking stderr")
				return
			default:
				buf, err := ioutil.ReadAll(stderr)
				if err != nil {
					r.cfg.Logger.Log(logger.Error, "could not read stderr", "error", err)
					return
				}

				if len(buf) != 0 {
					r.cfg.Logger.Log(logger.Error, "error from raspivid stderr", "error", string(buf))
					return
				}
			}
		}
	}()

	err = r.cmd.Start()
	if err != nil {
		return fmt.Errorf("could not start raspivid command: %w", err)
	}

	return nil
}

// Read implements io.Reader. Calling read before Start has been called will
// result in 0 bytes read and an error.
func (r *Raspivid) Read(p []byte) (int, error) {
	if r.out != nil {
		return r.out.Read(p)
	}
	return 0, errors.New("cannot read, raspivid has not started")
}

// Stop will terminate the raspivid process and close the output pipe.
func (r *Raspivid) Stop() error {
	close(r.done)
	if r.cmd == nil || r.cmd.Process == nil {
		return errors.New("raspivid process was never started")
	}
	err := r.cmd.Process.Kill()
	if err != nil {
		return fmt.Errorf("could not kill raspivid process: %w", err)
	}
	return r.out.Close()
}