Merged in get-pts-range (pull request #197)

container/mts: GetPTSRange func

Approved-by: Alan Noble <anoble@gmail.com>

container/mts/encoder.go

@@ -114,14 +114,17 @@ const (
 	Audio
 )
 
-// Time related constants.
+// Time-related constants.
 const (
 	// ptsOffset is the offset added to the clock to determine
 	// the current presentation timestamp.
 	ptsOffset = 700 * time.Millisecond
 
-	// pcrFreq is the base Program Clock Reference frequency.
-	pcrFreq = 90000 // Hz
+	// PCRFrequency is the base Program Clock Reference frequency in Hz.
+	PCRFrequency = 90000
+
+	// PTSFrequency is the presentation timestamp frequency in Hz.
+	PTSFrequency = 90000
 )
 
 // Encoder encapsulates properties of an MPEG-TS generator.
@@ -300,12 +303,12 @@ func (e *Encoder) tick() {
 
 // pts retuns the current presentation timestamp.
 func (e *Encoder) pts() uint64 {
-	return uint64((e.clock + e.ptsOffset).Seconds() * pcrFreq)
+	return uint64((e.clock + e.ptsOffset).Seconds() * PTSFrequency)
 }
 
 // pcr returns the current program clock reference.
 func (e *Encoder) pcr() uint64 {
-	return uint64(e.clock.Seconds() * pcrFreq)
+	return uint64(e.clock.Seconds() * PCRFrequency)
 }
 
 // ccFor returns the next continuity counter for pid.

container/mts/mpegts.go

@@ -26,6 +26,7 @@ LICENSE
   along with revid in gpl.txt.  If not, see [GNU licenses](http://www.gnu.org/licenses).
 */
 
+// Package mts provides MPEGT-TS (mts) encoding and related functions.
 package mts
 
 import (
@@ -33,6 +34,7 @@ import (
 	"fmt"
 
 	"github.com/Comcast/gots/packet"
+	"github.com/Comcast/gots/pes"
 )
 
 const PacketSize = 188
@@ -315,3 +317,47 @@ func DiscontinuityIndicator(f bool) Option {
 		p[DiscontinuityIndicatorIdx] |= DiscontinuityIndicatorMask & set
 	}
 }
+
+// GetPTSRange retreives the first and last PTS of an MPEGTS clip.
+func GetPTSRange(clip []byte, pid uint16) (pts [2]uint64, err error) {
+	// Find the first packet with PID pidType.
+	pkt, _, err := FindPid(clip, pid)
+	if err != nil {
+		return [2]uint64{}, err
+	}
+
+	// Get the payload of the packet, which will be the start of the PES packet.
+	var _pkt packet.Packet
+	copy(_pkt[:], pkt)
+	payload, err := packet.Payload(&_pkt)
+	if err != nil {
+		fmt.Printf("_pkt: %v\n", _pkt)
+		return [2]uint64{}, err
+	}
+
+	// Get the the first PTS from the PES header.
+	_pes, err := pes.NewPESHeader(payload)
+	if err != nil {
+		return [2]uint64{}, err
+	}
+	pts[0] = _pes.PTS()
+
+	// Get the final PTS searching from end of clip for access unit start.
+	for i := len(clip) - PacketSize; i >= 0; i -= PacketSize {
+		copy(_pkt[:], clip[i:i+PacketSize])
+		if packet.PayloadUnitStartIndicator(&_pkt) && uint16(_pkt.PID()) == pid {
+			payload, err = packet.Payload(&_pkt)
+			if err != nil {
+				return [2]uint64{}, err
+			}
+			_pes, err = pes.NewPESHeader(payload)
+			if err != nil {
+				return [2]uint64{}, err
+			}
+			pts[1] = _pes.PTS()
+			return
+		}
+	}
+
+	return [2]uint64{}, errors.New("could only find one access unit in mpegts clip")
+}

container/mts/mpegts_test.go

@@ -29,11 +29,134 @@ package mts
 
 import (
 	"bytes"
+	"math/rand"
 	"testing"
+	"time"
 
+	"bitbucket.org/ausocean/av/container/mts/pes"
+	"bitbucket.org/ausocean/av/container/mts/psi"
 	"github.com/Comcast/gots/packet"
 )
 
+// TestGetPTSRange checks that GetPTSRange can correctly get the first and last
+// PTS in an MPEGTS clip.
+func TestGetPTSRange(t *testing.T) {
+	const (
+		numOfFrames  = 20
+		maxFrameSize = 1000
+		minFrameSize = 100
+		rate         = 25                // fps
+		interval     = float64(1) / rate // s
+		ptsFreq      = 90000             // Hz
+	)
+
+	// Generate randomly sized data for each frame.
+	rand.Seed(time.Now().UnixNano())
+	frames := make([][]byte, numOfFrames)
+	for i := range frames {
+		size := rand.Intn(maxFrameSize-minFrameSize) + minFrameSize
+		frames[i] = make([]byte, size)
+	}
+
+	var clip bytes.Buffer
+
+	// Write the PSI first.
+	err := writePSI(&clip)
+	if err != nil {
+		t.Fatalf("did not expect error writing psi: %v", err)
+	}
+
+	// Now write frames.
+	var curTime float64
+	for _, frame := range frames {
+		nextPTS := curTime * ptsFreq
+
+		err = writeFrame(&clip, frame, uint64(nextPTS))
+		if err != nil {
+			t.Fatalf("did not expect error writing frame: %v", err)
+		}
+
+		curTime += interval
+	}
+
+	got, err := GetPTSRange(clip.Bytes(), videoPid)
+	if err != nil {
+		t.Fatalf("did not expect error getting PTS range: %v", err)
+	}
+
+	want := [2]uint64{0, uint64((numOfFrames - 1) * interval * ptsFreq)}
+	if got != want {
+		t.Errorf("did not get expected result.\n Got: %v\n Want: %v\n", got, want)
+	}
+}
+
+// writePSI is a helper function write the PSI found at the start of a clip.
+func writePSI(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
+	}
+
+	// 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
+}
+
+// writeFrame is a helper function used to form a PES packet from a frame, and
+// then fragment this across MPEGTS packets where they are then written to the
+// given buffer.
+func writeFrame(b *bytes.Buffer, frame []byte, pts uint64) error {
+	// Prepare PES data.
+	pesPkt := pes.Packet{
+		StreamID:     videoStreamID,
+		PDI:          hasPTS,
+		PTS:          pts,
+		Data:         frame,
+		HeaderLength: 5,
+	}
+	buf := pesPkt.Bytes(nil)
+
+	// Write PES data acroos MPEGTS packets.
+	pusi := true
+	for len(buf) != 0 {
+		pkt := Packet{
+			PUSI: pusi,
+			PID:  videoPid,
+			RAI:  pusi,
+			CC:   0,
+			AFC:  hasAdaptationField | hasPayload,
+			PCRF: pusi,
+		}
+		n := pkt.FillPayload(buf)
+		buf = buf[n:]
+
+		pusi = false
+		_, err := b.Write(pkt.Bytes(nil))
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // TestBytes checks that Packet.Bytes() correctly produces a []byte
 // representation of a Packet.
 func TestBytes(t *testing.T) {

container/mts/pes/pes.go

@@ -26,6 +26,8 @@ LICENSE
 
 package pes
 
+import "github.com/Comcast/gots"
+
 const MaxPesSize = 64 * 1 << 10
 
 /*
@@ -108,16 +110,11 @@ func (p *Packet) Bytes(buf []byte) []byte {
 			boolByte(p.ACIF)<<2 | boolByte(p.CRCF)<<1 | boolByte(p.EF)),
 		p.HeaderLength,
 	}...)
+
 	if p.PDI == byte(2) {
-		pts := 0x2100010001 | (p.PTS&0x1C0000000)<<3 | (p.PTS&0x3FFF8000)<<2 |
-			(p.PTS&0x7FFF)<<1
-		buf = append(buf, []byte{
-			byte((pts & 0xFF00000000) >> 32),
-			byte((pts & 0x00FF000000) >> 24),
-			byte((pts & 0x0000FF0000) >> 16),
-			byte((pts & 0x000000FF00) >> 8),
-			byte(pts & 0x00000000FF),
-		}...)
+		ptsIdx := len(buf)
+		buf = buf[:ptsIdx+5]
+		gots.InsertPTS(buf[ptsIdx:], p.PTS)
 	}
 	buf = append(buf, append(p.Stuff, p.Data...)...)
 	return buf