container/mts: added payload_test.go and fixed bugs in Extract

container/mts/meta/meta.go

@@ -216,7 +216,7 @@ func GetAll(d []byte) ([][2]string, error) {
 	return all, nil
 }
 
-// GetAllAsMap returns a map containging keys and values from a slice d containing
+// GetAllAsMap returns a map containing keys and values from a slice d containing
 // AusOcean metadata.
 func GetAllAsMap(d []byte) (map[string]string, error) {
 	err := checkMeta(d)

container/mts/payload.go

@@ -15,19 +15,25 @@ func Extract(p []byte) (Clip, error) {
 		return nil, errors.New("MTS clip is not of valid size")
 	}
 
-	// This will hold a copy of all the media in the MTS clip.
+	// This will hold a copy of all the media in the MPEG-TS clip.
 	buf := make([]byte, 0, l/PacketSize)
 
-	var clip Clip
+	var (
-	var meta map[string]string
+		clip        Clip              // The data that will be returned.
-	var err error
+		meta        map[string]string // Holds the most recently extracted meta.
-	var prev *Frame
+		lenOfFrame  int               // Len of current frame.
-	// Go through and get the number of frames, and also their size and allocate
+		dataLen     int               // Len of data from MPEG-TS packet.
-	// mem to clip.
+		curPTS      uint64            // Holds the current PTS.
+		curStreamID uint8             // Holds current StreamID (shouldn't change)
+		firstPUSI   = true            // Indicates that we have not yet received a PUSI.
+		err         error
+	)
+
+	// Go through the MPEGT-TS packets.
 	var pkt packet.Packet
 	for i := 0; i < l; i += PacketSize {
+		// We will use comcast/gots Packet type, so copy in.
 		copy(pkt[:], p[i:i+PacketSize])
-		idx := len(buf)
 
 		switch pkt.PID() {
 		case PatPid: // Do nothing.
@@ -36,8 +42,8 @@ func Extract(p []byte) (Clip, error) {
 			if err != nil {
 				return nil, err
 			}
-		default:
+		default: // Must be media.
-			// Get the MTS payload.
+			// Get the MPEG-TS payload.
 			payload, err := pkt.Payload()
 			if err != nil {
 				return nil, err
@@ -45,40 +51,46 @@ func Extract(p []byte) (Clip, error) {
 
 			// If PUSI is true then we know it's the start of a new frame, and we have
 			// a PES header in the MTS payload.
+
 			if pkt.PayloadUnitStartIndicator() {
 				_pes, err := pes.NewPESHeader(payload)
 				if err != nil {
 					return nil, err
 				}
-				pesData := _pes.Data()
-				lenOfFrame := len(pesData)
 
-				// If we have a previous Frame, then we need to adjust upper bound of
+				// Extract the PTS and ID, then add a new frame to the clip.
-				// it's media data.
+				curPTS = _pes.PTS()
-				if prev != nil {
+				curStreamID = _pes.StreamId()
-					prev.Media = prev.Media[:lenOfFrame]
+				clip = append(clip, Frame{
-				}
+					PTS:  curPTS,
+					ID:   curStreamID,
+					Meta: meta,
+				})
 
-				// Create a new frame.
+				// Append the data to the underlying buffer and get appended lenghth.
-				prev = &Frame{
+				buf = append(buf, _pes.Data()...)
-					Media: buf[idx:],
+				dataLen = len(_pes.Data())
-					PTS:   _pes.PTS(),
+
-					Meta:  meta,
+				// If we haven't hit the first PUSI, then we know we have a full frame
+				// and can add this data to the frame pertaining to the finish frame.
+				if !firstPUSI {
+					clip[len(clip)-2].Media = buf[:lenOfFrame]
+					buf = buf[lenOfFrame:]
+					lenOfFrame = 0
 				}
-				// Append media data to underlying media buffer.
+				firstPUSI = false
-				buf = append(buf, pesData...)
 			} else {
+				// We're not at the start of the frame, so we don't have a PES header.
+				// We can append the MPEG-TS data directly to the underlying buf.
+				dataLen = len(payload)
 				buf = append(buf, payload...)
 			}
-
+			lenOfFrame += dataLen
-			// And then append to the clip.
-			clip = append(clip, *prev)
 		}
 	}
-
+	// We're finished up with media frames, so give the final Frame it's data.
-	// Copy over data
+	clip[len(clip)-1].Media = buf[:lenOfFrame]
-
+	return clip, nil
-	return nil, nil
 }
 
 // Clip represents a clip of media, i.e. a sequence of media frames.
@@ -88,7 +100,7 @@ type Clip []Frame
 type Frame struct {
 	Media []byte            // Contains the media from the frame.
 	PTS   uint64            // PTS from PES packet (this gives time relative from start of stream).
-	ID    int8              // StreamID from the PES packet, identifying media codec.
+	ID    uint8             // StreamID from the PES packet, identifying media codec.
 	Meta  map[string]string // Contains metadata from PMT relevant to this frame.
 }
 

container/mts/payload_test.go

@@ -0,0 +1,137 @@
+package mts
+
+import (
+	"bytes"
+	"math/rand"
+	"reflect"
+	"strconv"
+	"testing"
+	"time"
+
+	"bitbucket.org/ausocean/av/container/mts/meta"
+	"bitbucket.org/ausocean/av/container/mts/psi"
+)
+
+// TestExtract checks that we can coorectly extract media, pts, id and meta from
+// an mpegts stream using Extract.
+func TestExtract(t *testing.T) {
+	Meta = meta.New()
+
+	const (
+		psiInterval  = 5                 // Write PSI at start and after every 5 frames.
+		numOfFrames  = 30                // Total number of frames to write.
+		maxFrameSize = 1000              // Max frame size to randomly generate.
+		minFrameSize = 100               // Min frame size to randomly generate.
+		rate         = 25                // Framerate (fps)
+		interval     = float64(1) / rate // Time interval between frames.
+		ptsFreq      = 90000             // Standard PTS frequency base.
+	)
+
+	// Generate randomly sized data for each frame and fill.
+	rand.Seed(time.Now().UnixNano())
+	frames := make([][]byte, numOfFrames)
+	for i := range frames {
+		size := rand.Intn(maxFrameSize-minFrameSize) + minFrameSize
+		frames[i] = make([]byte, size)
+		for j := 0; j < len(frames[i]); j++ {
+			frames[i][j] = byte(j)
+		}
+	}
+
+	var (
+		clip bytes.Buffer // This will hold the MPEG-TS data.
+		want Clip         // This is the Clip that we should get.
+		err  error
+	)
+
+	// Now write frames.
+	var curTime float64
+	for i, frame := range frames {
+		// Check to see if it's time to write another lot of PSI.
+		if i%psiInterval == 0 && i != len(frames)-1 {
+			// We'll add the frame number as meta.
+			Meta.Add("frameNum", strconv.Itoa(i))
+
+			err = writePSIWithMeta(&clip)
+			if err != nil {
+				t.Fatalf("did not expect error writing psi: %v", err)
+			}
+		}
+		nextPTS := uint64(curTime * ptsFreq)
+
+		err = writeFrame(&clip, frame, uint64(nextPTS))
+		if err != nil {
+			t.Fatalf("did not expect error writing frame: %v", err)
+		}
+
+		curTime += interval
+
+		// Need the meta map for the new expected Frame.
+		metaMap, err := meta.GetAllAsMap(Meta.Encode())
+		if err != nil {
+			t.Fatalf("did not expect error getting meta map: %v", err)
+		}
+
+		// Create an equivalent Frame and append to our Clip want.
+		want = append(want, Frame{
+			Media: frame,
+			PTS:   nextPTS,
+			ID:    H264ID,
+			Meta:  metaMap,
+		})
+	}
+
+	// Now use Extract to get frames from clip.
+	got, err := Extract(clip.Bytes())
+	if err != nil {
+		t.Fatalf("did not expect error using Extract. Err: %v", err)
+	}
+
+	// Check length of got and want.
+	if len(want) != len(got) {
+		t.Fatalf("did not get expected length for got.\nGot: %v\n, Want: %v\n", len(got), len(want))
+	}
+
+	// Check frames individually.
+	for i, frame := range want {
+		if !reflect.DeepEqual(frame, got[i]) {
+			t.Fatalf("did not get expected result.\nGot: %v\n, Want: %v\n", got[i], frame)
+		}
+	}
+}
+
+// writePSIWithMeta writes PSI to b with updated metadata.
+func writePSIWithMeta(b *bytes.Buffer) error {
+	// Write PAT.
+	pat := Packet{
+		PUSI:    true,
+		PID:     PatPid,
+		CC:      0,
+		AFC:     HasPayload,
+		Payload: psi.AddPadding(patTable),
+	}
+	_, err := b.Write(pat.Bytes(nil))
+	if err != nil {
+		return err
+	}
+
+	// Update the meta in the pmt table.
+	pmtTable, err = updateMeta(pmtTable)
+	if err != nil {
+		return err
+	}
+
+	// Write PMT.
+	pmt := Packet{
+		PUSI:    true,
+		PID:     PmtPid,
+		CC:      0,
+		AFC:     HasPayload,
+		Payload: psi.AddPadding(pmtTable),
+	}
+	_, err = b.Write(pmt.Bytes(nil))
+	if err != nil {
+		return err
+	}
+	return nil
+}