av/container/mts/payload.go

package mts

import (
	"errors"

	"github.com/Comcast/gots/packet"
	"github.com/Comcast/gots/pes"
)

// Extract extracts the media from an MPEG-TS clip given by p.
func Extract(p []byte) (Clip, error) {
	l := len(p)
	// Check that clip is divisible by 188, i.e. contains a series of full MPEG-TS clips.
	if l%PacketSize != 0 {
		return nil, errors.New("MTS clip is not of valid size")
	}

	// This will hold a copy of all the media in the MPEG-TS clip.
	buf := make([]byte, 0, l/PacketSize)

	var (
		clip        Clip              // The data that will be returned.
		meta        map[string]string // Holds the most recently extracted meta.
		lenOfFrame  int               // Len of current frame.
		dataLen     int               // Len of data from MPEG-TS packet.
		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])

		switch pkt.PID() {
		case PatPid: // Do nothing.
		case PmtPid:
			meta, err = ExtractMeta(pkt[:])
			if err != nil {
				return nil, err
			}
		default: // Must be media.
			// Get the MPEG-TS payload.
			payload, err := pkt.Payload()
			if err != nil {
				return nil, err
			}

			// 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
				}

				// Extract the PTS and ID, then add a new frame to the clip.
				curPTS = _pes.PTS()
				curStreamID = _pes.StreamId()
				clip = append(clip, Frame{
					PTS:  curPTS,
					ID:   curStreamID,
					Meta: meta,
				})

				// Append the data to the underlying buffer and get appended lenghth.
				buf = append(buf, _pes.Data()...)
				dataLen = len(_pes.Data())

				// 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
				}
				firstPUSI = false
			} 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
		}
	}
	// We're finished up with media frames, so give the final Frame it's data.
	clip[len(clip)-1].Media = buf[:lenOfFrame]
	return clip, nil
}

// Clip represents a clip of media, i.e. a sequence of media frames.
type Clip []Frame

// Frame describes a media frame that may be extracted from a PES packet.
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    uint8             // StreamID from the PES packet, identifying media codec.
	Meta  map[string]string // Contains metadata from PMT relevant to this frame.
}

// Bytes returns the concatentated media bytes from each frame in the Clip c.
func (c *Clip) Bytes() []byte {
	return nil
}