av/generator/FLVGenerator.go

/*
NAME
  RtpToTsConverter.go - provides utilities for the conversion of Rtp packets
  to equivalent MpegTs packets.

DESCRIPTION
  See Readme.md

AUTHOR
  Saxon Nelson-Milton <saxon.milton@gmail.com>

LICENSE
  RtpToTsConverter.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 [GNU licenses](http://www.gnu.org/licenses).
*/
package generator

import (
	"../flv"
)

const (
	inputChanLength = 1000
	outputChanLength = 1000
)

type flvGenerator struct {
	fps							uint
  inputChan chan []byte
  outputChan chan []byte
	audioFlag bool
	videoFlag bool
	lastTagSize int
	currentTimestamp uint32
  header flv.Header
}

func (g *flvGenerator)GetInputChan() chan []byte {
	return g.inputChan
}

func (g *flvGenerator)GetOutputChan() chan []byte {
	return g.outputChan
}

func NewFlvGenerator(audio bool, video bool, fps uint) (g *flvGenerator) {
	g = new(flvGenerator)
	g.fps = fps
	g.audioFlag = audio
	g.videoFlag = video
  g.currentTimestamp = 0
	g.lastTagSize = 0
	g.inputChan = make(chan []byte, inputChanLength)
	g.outputChan = make(chan []byte, outputChanLength)
	return
}

func (g *flvGenerator) Start(){
	go g.generate()
}

func (g *flvGenerator) GenHeader(){
  header := flv.Header{
		AudioFlag: g.audioFlag,
		VideoFlag: g.videoFlag,
  }
  g.outputChan <- header.ToByteSlice()
}

func (g *flvGenerator) getNextTimestamp() (timestamp uint32){
  timestamp = g.currentTimestamp
  g.currentTimestamp += uint32(1000) / uint32(g.fps)
  return
}

func (g *flvGenerator) ResetTimestamp() {
  g.currentTimestamp = 0
}

func (g *flvGenerator) generate() {
  g.GenHeader()
  for {
    select {
    case videoFrame := <-g.inputChan:
			tag := flv.VideoTag{
				PrevTagSize: uint32(g.lastTagSize),
				TagType: uint8(flv.VideoTagType),
				DataSize: uint32(len(videoFrame)) + flv.DataHeaderLength,
				Timestamp: g.getNextTimestamp(),
				TimestampExtended: flv.NoTimestampExtension,
				FrameType: flv.KeyFrameType,
				Codec: flv.H264,
				PacketType: flv.AVCNALU,
				CompositionTime: 0,
				Data: videoFrame,
			}
			tagAsByteSlice := tag.ToByteSlice()
			g.lastTagSize = len(tagAsByteSlice)
			g.outputChan<-tagAsByteSlice
    }
  }
}