av/revid/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 revid

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

	"bitbucket.org/ausocean/av/codec/codecutil"
	"bitbucket.org/ausocean/utils/logger"
)

// Raspivid AVDevice configuration defaults.
const (
	raspividDefaultCodec            = codecutil.H264
	raspividDefaultRotation         = 0
	raspividDefaultWidth            = 1280
	raspividDefaultHeight           = 720
	raspividDefaultBrightness       = 50
	raspividDefaultSaturation       = 0
	raspividDefaultExposure         = "auto"
	raspividDefaultAutoWhiteBalance = "auto"
	raspividDefaultMinFrames        = 100
	raspividDefaultQuantization     = 30
	raspividDefaultBitrate          = 48000
	raspividDefaultFramerate        = 25
)

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")
)

type Raspivid struct {
	cfg Config
	cmd *exec.Cmd
	out io.ReadCloser
}

type multiError []error

func (me multiError) Error() string {
	return fmt.Sprintf("%v", me)
}

func (r *Raspivid) Set(c Config) error {
	var errs []error
	switch c.InputCodec {
	case codecutil.H264, codecutil.MJPEG:
	default:
		c.InputCodec = raspividDefaultCodec
		errs = append(errs, errBadCodec)
	}

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

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

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

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

	if c.VBR {
		c.Bitrate = 0
		if c.Quantization < 10 || c.Quantization > 40 {
			errs = append(errs, errBadQuantization)
			c.Quantization = raspividDefaultQuantization
		}
	} else {
		c.Quantization = 0
		if c.Bitrate <= 0 {
			errs = append(errs, errBadBitrate)
			c.Bitrate = raspividDefaultBitrate
		}
	}

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

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

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

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

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

	r.cfg = c
	return multiError(errs)
}

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.FlipHorizontal {
		args = append(args, "--hflip")
	}

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

	switch r.cfg.InputCodec {
	default:
		return fmt.Errorf("revid: invalid input codec: %v", r.cfg.InputCodec)
	case codecutil.H264:
		args = append(args,
			"--codec", "H264",
			"--inline",
			"--intra", fmt.Sprint(r.cfg.MinFrames),
		)
		if r.cfg.VBR {
			args = append(args, "-qp", fmt.Sprint(r.cfg.Quantization))
		}
	case codecutil.MJPEG:
		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)
	}

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

	return nil
}

func (r *Raspivid) Read(p []byte) (int, error) {
	return r.out.Read(p)
}

func (r *Raspivid) Stop() error {
	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 nil
}