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"
	_"fmt"
	"time"
)

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 isKeyFrame(frame []byte) bool {
	byteChannel := make(chan byte, len(frame))
	for i := range frame {
		byteChannel <- frame[i]
	}
	for len(byteChannel) >= 5{
    aByte := <-byteChannel
		for i:=1; aByte == 0x00 && i != 4; i++ {
			aByte = <-byteChannel
			if ( aByte == 0x01 && i == 2 ) || ( aByte == 0x01 && i == 3 ) {
				aByte = <-byteChannel
				nalType := aByte & 0x1F
				switch nalType {
				case 1:
					return false
				case 5:
					return true
				case 6:
					return true
				}
			}
		}
	}
	return false
}

func isSequenceHeader(frame []byte) bool {
	byteChannel := make(chan byte, len(frame))
	for i := range frame {
		byteChannel <- frame[i]
	}
	for len(byteChannel) >= 5{
    aByte := <-byteChannel
		for i:=1; aByte == 0x00 && i != 4; i++ {
			aByte = <-byteChannel
			if ( aByte == 0x01 && i == 2 ) || ( aByte == 0x01 && i == 3 ) {
				aByte = <-byteChannel
				nalType := aByte & 0x1F
				switch nalType {
				case 1:
					return false
				case 5:
					return false
				case 6:
					return true
				case 7:
					return true
				case 8:
					return true
				}
			}
		}
	}
	return false
}

func (g *flvGenerator) generate() {
  g.GenHeader()
  for {
    select {
    case videoFrame := <-g.inputChan:
			var frameType byte
			if isKeyFrame(videoFrame) {
				frameType = flv.KeyFrameType
			} else {
				frameType = flv.InterFrameType
			}
			var packetType byte
			if isSequenceHeader(videoFrame){
				packetType = flv.SequenceHeader
			} else {
				packetType = flv.AVCNALU

			}
			timeStamp := g.getNextTimestamp()
			videoTag := flv.VideoTag{
				PrevTagSize: uint32(g.lastTagSize),
				TagType: uint8(flv.VideoTagType),
				DataSize: uint32(len(videoFrame)) + flv.DataHeaderLength,
				Timestamp: timeStamp,
				TimestampExtended: flv.NoTimestampExtension,
				FrameType: frameType,
				Codec: flv.H264,
				PacketType: packetType,
				CompositionTime: 0,
				Data: videoFrame,
			}
			videoTagAsByteSlice := videoTag.ToByteSlice()
			g.lastTagSize = len(videoTagAsByteSlice)
			g.outputChan<-videoTagAsByteSlice

			soundData := make([]byte, 10)
			for i := range soundData {
				if i == 0 {
					soundData[i] = 1
				} else {
					soundData[i] = 0
				}
			}
			audioTag := flv.AudioTag{
				PrevTagSize: uint32(g.lastTagSize),
				TagType: uint8(flv.AudioTagType),
				DataSize: 7,
				Timestamp: timeStamp,
				TimestampExtended: flv.NoTimestampExtension,
				SoundFormat: flv.AACAudioFormat,
				SoundRate: 3,
				SoundSize: true,
				SoundType: true,
				Data: []byte{0x00,0x12,0x08,0x56,0xe5,0x00},
			}
			audioTagAsByteSlice := audioTag.ToByteSlice()
			g.lastTagSize = len(audioTagAsByteSlice)
			g.outputChan<-audioTagAsByteSlice

		time.Sleep(60*time.Millisecond)
    }
  }
}