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
	audioSize = 18
	videoHeaderSize = 16
	interFrameCode = 1
	keyFrameCode = 5
	sequenceCode = 6
)

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 interFrameCode:
					return false
				case keyFrameCode:
					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()

			if g.videoFlag {
				videoTag := flv.VideoTag{
					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,
					PrevTagSize: uint32(videoHeaderSize+len(videoFrame)),
				}
				videoTagAsByteSlice := videoTag.ToByteSlice()
				g.outputChan<-videoTagAsByteSlice
			}

			// TODO: Create some more constants
 
			if g.audioFlag {
				audioTag := flv.AudioTag{
					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},
					PrevTagSize: uint32(audioSize),
				}
				audioTagAsByteSlice := audioTag.ToByteSlice()
				g.outputChan<-audioTagAsByteSlice

				audioTag = flv.AudioTag{
					TagType: uint8(flv.AudioTagType),
					DataSize: 21,
					Timestamp: timeStamp,
					TimestampExtended: flv.NoTimestampExtension,
					SoundFormat: flv.AACAudioFormat,
					SoundRate: 3,
					SoundSize: true,
					SoundType: true,
					Data: []byte{0x01,0xdc,0x00,0x4c,0x61,0x76,0x63,0x35,0x38,0x2e,0x36,
						0x2e,0x31,0x30,0x32,0x00,0x02,0x30,0x40,0x0e,},
					PrevTagSize: uint32(22),
				}
				audioTagAsByteSlice = audioTag.ToByteSlice()
				g.outputChan<-audioTagAsByteSlice
			}
    }
  }
}