av/generator/flv_generator.go

/*
NAME
  flv_generator.go

DESCRIPTION
  See Readme.md

AUTHOR
  Saxon Nelson-Milton <saxon@ausocean.org>

LICENSE
  flv_generator.go is Copyright (C) 2017 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 generator

import (
	"time"

	"bitbucket.org/ausocean/av/flv"
)

const (
	inputChanLength  = 500
	outputChanLength = 500
	audioSize        = 18
	videoHeaderSize  = 16
)

// Data representing silent audio (required for youtube)
var (
	dummyAudioTag1Data = []byte{
		0x00, 0x12, 0x08, 0x56, 0xe5, 0x00,
	}

	dummyAudioTag2Data = []byte{
		0x01, 0xdc, 0x00, 0x4c, 0x61, 0x76, 0x63, 0x35, 0x38,
		0x2e, 0x36, 0x2e, 0x31, 0x30, 0x32, 0x00, 0x02, 0x30,
		0x40, 0x0e,
	}
)

// flvGenerator provides properties required for the generation of flv video
// from raw video data
type flvGenerator struct {
	fps          int
	inputChan    chan []byte
	outputChan   chan []byte
	audio        bool
	video        bool
	lastTagSize  int
	header       flv.Header
	startTime    time.Time
	firstTag     bool
	isGenerating bool
}

// InputChan returns the input channel to the generator. This is where the
// raw data frames are entered into the generator
func (g *flvGenerator) InputChan() chan []byte {
	return g.inputChan
}

// OutputChan retuns the output chan of the generator - this is where the
// flv packets (more specifically tags) are outputted.
func (g *flvGenerator) OutputChan() <-chan []byte {
	return g.outputChan
}

// NewFlvGenerator retuns an instance of the flvGenerator struct
func NewFlvGenerator(audio, video bool, fps int) *flvGenerator {
	return &flvGenerator{
		fps:        fps,
		audio:      audio,
		video:      video,
		inputChan:  make(chan []byte, inputChanLength),
		outputChan: make(chan []byte, outputChanLength),
		firstTag:   true,
	}
}

// Start begins the generation routine - i.e. if raw data is given to the input
// channel flv tags will be produced and available from the output channel.
func (g *flvGenerator) Start() {
	g.isGenerating = true
	go g.generate()
}

func (g *flvGenerator) Stop() {
	g.isGenerating = false
}

// GenHeader generates the flv header and sends it down the output chan for use
// This will generally be called once at the start of file writing/transmission.
func (g *flvGenerator) GenHeader() {
	header := flv.Header{
		HasAudio: g.audio,
		HasVideo: g.video,
	}
	g.outputChan <- header.Bytes()
}

// getNextTimestamp generates and returns the next timestamp based on current time
func (g *flvGenerator) getNextTimestamp() (timestamp uint32) {
	if g.firstTag {
		g.startTime = time.Now()
		g.firstTag = false
		return 0
	}
	return uint32(time.Now().Sub(g.startTime).Seconds() * float64(1000))
}

// http://www.itu.int/rec/dologin_pub.asp?lang=e&id=T-REC-H.264-200305-S!!PDF-E&type=items
// Table 7-1 NAL unit type codes
const (
	nonIdrPic   = 1
	idrPic      = 5
	suppEnhInf  = 6
	seqParamSet = 7
	paramSet    = 8
)

// isKeyFrame returns true if the passed frame data represents that of a keyframe
// FIXME(kortschak): Clarify and document the logic of this functions.
func isKeyFrame(frame []byte) bool {
	sc := frameScanner{buf: frame}
	for {
		b, ok := sc.readByte()
		if !ok {
			return false
		}
		for i := 1; b == 0x00 && i < 4; i++ {
			b, ok = sc.readByte()
			if !ok {
				return false
			}
			if b != 0x01 || (i != 3 && i != 2) {
				continue
			}

			b, ok = sc.readByte()
			if !ok {
				return false
			}
			switch nalTyp := b & 0x1f; nalTyp {
			case idrPic, suppEnhInf:
				return true
			case nonIdrPic:
				return false
			}
		}
	}
	return false
}

// isSequenceHeader returns true if the passed frame data represents that of a
// a sequence header.
// FIXME(kortschak): Clarify and document the logic of this functions.
func isSequenceHeader(frame []byte) bool {
	sc := frameScanner{buf: frame}
	for {
		b, ok := sc.readByte()
		if !ok {
			return false
		}
		for i := 1; b == 0x00 && i != 4; i++ {
			b, ok = sc.readByte()
			if !ok {
				return false
			}
			if b != 0x01 || (i != 2 && i != 3) {
				continue
			}

			b, ok = sc.readByte()
			if !ok {
				return false
			}
			switch nalTyp := b & 0x1f; nalTyp {
			case suppEnhInf, seqParamSet, paramSet:
				return true
			case nonIdrPic, idrPic:
				return false
			}
		}
	}
}

type frameScanner struct {
	off int
	buf []byte
}

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
}

// generate takes in raw video data from the input chan and packetises it into
// flv tags, which are then passed to the output channel.
func (g *flvGenerator) generate() {
	g.GenHeader()
	var frameType byte
	var packetType byte
	for g.isGenerating {
		select {
		default:
			time.Sleep(time.Duration(5) * time.Millisecond)
		case frame := <-g.inputChan:
			timeStamp := g.getNextTimestamp()
			// Do we have video to send off?
			if g.video {
				if isKeyFrame(frame) {
					frameType = flv.KeyFrameType
				} else {
					frameType = flv.InterFrameType
				}
				if isSequenceHeader(frame) {
					packetType = flv.SequenceHeader
				} else {
					packetType = flv.AVCNALU
				}

				tag := flv.VideoTag{
					TagType:           uint8(flv.VideoTagType),
					DataSize:          uint32(len(frame)) + flv.DataHeaderLength,
					Timestamp:         timeStamp,
					TimestampExtended: flv.NoTimestampExtension,
					FrameType:         frameType,
					Codec:             flv.H264,
					PacketType:        packetType,
					CompositionTime:   0,
					Data:              frame,
					PrevTagSize:       uint32(videoHeaderSize + len(frame)),
				}
				g.outputChan <- tag.Bytes()
			}
			// Do we even have some audio to send off ?
			if g.audio {
				// Not sure why but we need two audio tags for dummy silent audio
				// TODO: create constants or SoundSize and SoundType parameters
				tag := flv.AudioTag{
					TagType:           uint8(flv.AudioTagType),
					DataSize:          7,
					Timestamp:         timeStamp,
					TimestampExtended: flv.NoTimestampExtension,
					SoundFormat:       flv.AACAudioFormat,
					SoundRate:         3,
					SoundSize:         true,
					SoundType:         true,
					Data:              dummyAudioTag1Data,
					PrevTagSize:       uint32(audioSize),
				}
				g.outputChan <- tag.Bytes()

				tag = flv.AudioTag{
					TagType:           uint8(flv.AudioTagType),
					DataSize:          21,
					Timestamp:         timeStamp,
					TimestampExtended: flv.NoTimestampExtension,
					SoundFormat:       flv.AACAudioFormat,
					SoundRate:         3,
					SoundSize:         true,
					SoundType:         true,
					Data:              dummyAudioTag2Data,
					PrevTagSize:       uint32(22),
				}
				g.outputChan <- tag.Bytes()
			}
		}
	}
}