mts: Added audio mts encoding and test

container/mts/audio_test.go

@@ -0,0 +1,109 @@
+/*
+NAME
+  audio_test.go
+
+DESCRIPTION
+  See Readme.md
+
+AUTHOR
+  Trek Hopton <trek@ausocean.org>
+
+LICENSE
+  audio_test.go is Copyright (C) 2017-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
+  along with revid in gpl.txt. If not, see http://www.gnu.org/licenses.
+*/
+
+package mts
+
+import (
+	"bytes"
+	"io/ioutil"
+	"log"
+	"testing"
+
+	"bitbucket.org/ausocean/av/container/mts/meta"
+	"github.com/Comcast/gots/packet"
+	"github.com/Comcast/gots/pes"
+)
+
+// TestEncodePcm tests the mpegts encoder's ability to encode pcm audio data.
+// It reads and encodes input pcm data into mpegts, then decodes the mpegts compares the result to the input pcm.
+func TestEncodePcm(t *testing.T) {
+	Meta = meta.New()
+
+	var b []byte
+	buf := bytes.NewBuffer(b)
+	sampleRate := 48000
+	sampleSize := 2
+	writeSize := 16000
+	writeFreq := float64(sampleRate*sampleSize) / float64(writeSize)
+	e := NewEncoder(buf, writeFreq, Audio)
+
+	inPath := "../../../test/test-data/av/input/sweep_400Hz_20000Hz_-3dBFS_5s_48khz.pcm"
+	inPcm, err := ioutil.ReadFile(inPath)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	// Encode pcm to mts and get the resulting bytes.
+	_, err = e.Write(inPcm)
+	if err != nil {
+		log.Fatal(err)
+	}
+	clip := buf.Bytes()
+
+	// Decode the mts packets to extract the original data
+	var pkt packet.Packet
+	pesPacket := make([]byte, 0, writeSize)
+	got := make([]byte, 0, len(inPcm))
+	i := 0
+	for i+PacketSize <= len(clip) {
+		copy(pkt[:], clip[i:i+PacketSize])
+		if pkt.PID() == audioPid {
+			if pkt.PayloadUnitStartIndicator() {
+				payload, err := pkt.Payload()
+				if err != nil {
+					t.Fatalf("Unexpected err: %v\n", err)
+				}
+				pesPacket = append(pesPacket, payload...)
+				i += PacketSize
+				first := true
+				for (first || !pkt.PayloadUnitStartIndicator()) && i+PacketSize <= len(clip) {
+					first = false
+					copy(pkt[:], clip[i:i+PacketSize])
+					payload, err := pkt.Payload()
+					if err != nil {
+						t.Fatalf("Unexpected err: %v\n", err)
+					}
+
+					pesPacket = append(pesPacket, payload...)
+					i += PacketSize
+				}
+			}
+			pesHeader, err := pes.NewPESHeader(pesPacket)
+			if err != nil {
+				t.Fatalf("Unexpected err: %v\n", err)
+			}
+			got = append(got, pesHeader.Data()...)
+		} else {
+			i += PacketSize
+		}
+	}
+
+	// Compare encoded data with original data.
+	if !bytes.Equal(got, inPcm) {
+		t.Error("Error, unexpected output")
+	}
+}

container/mts/encoder.go

@@ -105,7 +105,16 @@ const (
 	patPid        = 0
 	pmtPid        = 4096
 	videoPid      = 256
-	streamID = 0xe0 // First video stream ID.
+	audioPid      = 210
+	videoStreamID = 0xe0 // First video stream ID.
+	audioStreamID = 0xc0 // First audio stream ID.
+)
+
+// Video and Audio constants are used to communicate which media type will be encoded when creating a
+// new encoder with NewEncoder.
+const (
+	Video = iota
+	Audio
 )
 
 // Time related constants.
@@ -124,7 +133,7 @@ type Encoder struct {
 
 	clock       time.Duration
 	lastTime    time.Time
-	frameInterval time.Duration
+	writePeriod time.Duration
 	ptsOffset   time.Duration
 	tsSpace     [PacketSize]byte
 	pesSpace    [pes.MaxPesSize]byte
@@ -134,26 +143,42 @@ type Encoder struct {
 	timeBasedPsi bool
 	pktCount     int
 	psiSendCount int
+	mediaPid     int
+	streamID     byte
 
 	psiLastTime time.Time
 }
 
 // NewEncoder returns an Encoder with the specified frame rate.
-func NewEncoder(dst io.Writer, fps float64) *Encoder {
+func NewEncoder(dst io.Writer, writeFreq float64, mediaType int) *Encoder {
+	var mPid int
+	var sid byte
+	switch mediaType {
+	case Audio:
+		mPid = audioPid
+		sid = audioStreamID
+	case Video:
+		mPid = videoPid
+		sid = videoStreamID
+	}
+
 	return &Encoder{
 		dst: dst,
 
-		frameInterval: time.Duration(float64(time.Second) / fps),
+		writePeriod: time.Duration(float64(time.Second) / writeFreq),
 		ptsOffset:   ptsOffset,
 
 		timeBasedPsi: true,
 
 		pktCount: 8,
 
+		mediaPid: mPid,
+		streamID: sid,
+
 		continuity: map[int]byte{
 			patPid: 0,
 			pmtPid: 0,
-			videoPid: 0,
+			mPid:   0,
 		},
 	}
 }
@@ -180,7 +205,7 @@ func (e *Encoder) TimeBasedPsi(b bool, sendCount int) {
 
 // Write implements io.Writer. Write takes raw h264 and encodes into mpegts,
 // then sending it to the encoder's io.Writer destination.
-func (e *Encoder) Write(nalu []byte) (int, error) {
+func (e *Encoder) Write(data []byte) (int, error) {
 	now := time.Now()
 	if (e.timeBasedPsi && (now.Sub(e.psiLastTime) > psiInterval)) || (!e.timeBasedPsi && (e.pktCount >= e.psiSendCount)) {
 		e.pktCount = 0
@@ -193,21 +218,22 @@ func (e *Encoder) Write(nalu []byte) (int, error) {
 
 	// Prepare PES data.
 	pesPkt := pes.Packet{
-		StreamID:     streamID,
+		StreamID:     e.streamID,
 		PDI:          hasPTS,
 		PTS:          e.pts(),
-		Data:         nalu,
+		Data:         data,
 		HeaderLength: 5,
 	}
+
 	buf := pesPkt.Bytes(e.pesSpace[:pes.MaxPesSize])
 
 	pusi := true
 	for len(buf) != 0 {
 		pkt := Packet{
 			PUSI: pusi,
-			PID:  videoPid,
+			PID:  uint16(e.mediaPid),
 			RAI:  pusi,
-			CC:   e.ccFor(videoPid),
+			CC:   e.ccFor(e.mediaPid),
 			AFC:  hasAdaptationField | hasPayload,
 			PCRF: pusi,
 		}
@@ -222,14 +248,14 @@ func (e *Encoder) Write(nalu []byte) (int, error) {
 		}
 		_, err := e.dst.Write(pkt.Bytes(e.tsSpace[:PacketSize]))
 		if err != nil {
-			return len(nalu), err
+			return len(data), err
 		}
 		e.pktCount++
 	}
 
 	e.tick()
 
-	return len(nalu), nil
+	return len(data), nil
 }
 
 // writePSI creates mpegts with pat and pmt tables - with pmt table having updated
@@ -271,7 +297,7 @@ func (e *Encoder) writePSI() error {
 
 // tick advances the clock one frame interval.
 func (e *Encoder) tick() {
-	e.clock += e.frameInterval
+	e.clock += e.writePeriod
 }
 
 // pts retuns the current presentation timestamp.

container/mts/metaEncode_test.go

@@ -49,7 +49,7 @@ func TestMetaEncode1(t *testing.T) {
 	Meta = meta.New()
 	var b []byte
 	buf := bytes.NewBuffer(b)
-	e := NewEncoder(buf, fps)
+	e := NewEncoder(buf, fps, Video)
 	Meta.Add("ts", "12345678")
 	if err := e.writePSI(); err != nil {
 		t.Errorf(errUnexpectedErr, err.Error())
@@ -78,7 +78,7 @@ func TestMetaEncode2(t *testing.T) {
 	Meta = meta.New()
 	var b []byte
 	buf := bytes.NewBuffer(b)
-	e := NewEncoder(buf, fps)
+	e := NewEncoder(buf, fps, Video)
 	Meta.Add("ts", "12345678")
 	Meta.Add("loc", "1234,4321,1234")
 	if err := e.writePSI(); err != nil {

container/mts/mpegts.go

@@ -74,7 +74,7 @@ const (
 )
 
 /*
-The below data struct encapsulates the fields of an MPEG-TS packet. Below is
+Packet encapsulates the fields of an MPEG-TS packet. Below is
 the formatting of an MPEG-TS packet for reference!
 
 													MPEG-TS Packet Formatting
@@ -196,7 +196,11 @@ func FindPid(d []byte, pid uint16) (pkt []byte, i int, err error) {
 func (p *Packet) FillPayload(data []byte) int {
 	currentPktLen := 6 + asInt(p.PCRF)*6 + asInt(p.OPCRF)*6 +
 		asInt(p.SPF)*1 + asInt(p.TPDF)*1 + len(p.TPD)
+	if len(data) > PayloadSize-currentPktLen {
 		p.Payload = make([]byte, PayloadSize-currentPktLen)
+	} else {
+		p.Payload = make([]byte, len(data))
+	}
 	return copy(p.Payload, data)
 }
 
@@ -214,7 +218,7 @@ func asByte(b bool) byte {
 	return 0
 }
 
-// ToByteSlice interprets the fields of the ts packet instance and outputs a
+// Bytes interprets the fields of the ts packet instance and outputs a
 // corresponding byte slice
 func (p *Packet) Bytes(buf []byte) []byte {
 	if buf == nil || cap(buf) != PacketSize {

container/mts/pes/pes.go

@@ -26,7 +26,7 @@ LICENSE
 
 package pes
 
-const MaxPesSize = 10000
+const MaxPesSize = 65536
 
 /*
 The below data struct encapsulates the fields of an PES packet. Below is