mpegts: support reading and writing MPEG-1/2 video (#76)

bluenviron · Sep 13, 2023 · cfb12a8 · cfb12a8
1 parent 93f9c90
commit cfb12a8
Show file tree

Hide file tree

Showing 7 changed files with 107 additions and 349 deletions.
diff --git a/pkg/formats/mpegts/codec_mpeg1_video.go b/pkg/formats/mpegts/codec_mpeg1_video.go
@@ -0,0 +1,22 @@
+package mpegts
+
+import (
+	"github.com/asticode/go-astits"
+)
+
+// CodecMPEG1Video is a MPEG-1/2 Video codec.
+type CodecMPEG1Video struct{}
+
+func (c CodecMPEG1Video) marshal(pid uint16) (*astits.PMTElementaryStream, error) {
+	return &astits.PMTElementaryStream{
+		ElementaryPID:               pid,
+		ElementaryStreamDescriptors: nil,
+		// we use MPEG-2 to notify readers that video can be either MPEG-1 or MPEG-2
+		StreamType: astits.StreamTypeMPEG2Video,
+	}, nil
+}
+
+// IsVideo implements Codec.
+func (CodecMPEG1Video) IsVideo() bool {
+	return true
+}
diff --git a/pkg/formats/mpegts/reader.go b/pkg/formats/mpegts/reader.go
@@ -18,8 +18,8 @@ type ReaderOnDecodeErrorFunc func(err error)
 // ReaderOnDataH26xFunc is the prototype of the callback passed to OnDataH26x.
 type ReaderOnDataH26xFunc func(pts int64, dts int64, au [][]byte) error
 
-// ReaderOnDataMPEG4VideoFunc is the prototype of the callback passed to OnDataMPEG4Video.
-type ReaderOnDataMPEG4VideoFunc func(pts int64, frame []byte) error
+// ReaderOnDataMPEGxVideoFunc is the prototype of the callback passed to OnDataMPEGxVideo.
+type ReaderOnDataMPEGxVideoFunc func(pts int64, frame []byte) error
 
 // ReaderOnDataOpusFunc is the prototype of the callback passed to OnDataOpus.
 type ReaderOnDataOpusFunc func(pts int64, packets [][]byte) error
@@ -121,14 +121,9 @@ func (r *Reader) OnDataH26x(track *Track, cb ReaderOnDataH26xFunc) {
 	}
 }
 
-// OnDataMPEG4Video sets a callback that is called when data from an MPEG-4 Video track is received.
-func (r *Reader) OnDataMPEG4Video(track *Track, cb ReaderOnDataMPEG4VideoFunc) {
+// OnDataMPEGxVideo sets a callback that is called when data from an MPEG-1/2/4 Video track is received.
+func (r *Reader) OnDataMPEGxVideo(track *Track, cb ReaderOnDataMPEGxVideoFunc) {
 	r.onData[track.PID] = func(pts int64, dts int64, data []byte) error {
-		if pts != dts {
-			r.onDecodeError(fmt.Errorf("PTS is not equal to DTS"))
-			return nil
-		}
-
 		return cb(pts, data)
 	}
 }

diff --git a/pkg/formats/mpegts/reader_test.go b/pkg/formats/mpegts/reader_test.go
@@ -297,6 +297,66 @@ var casesReadWriter = []struct {
 			},
 		},
 	},
+	{
+		"mpeg-1 video",
+		&Track{
+			PID:   257,
+			Codec: &CodecMPEG1Video{},
+		},
+		[]sample{
+			{
+				30 * 90000,
+				30 * 90000,
+				[][]byte{{0, 0, 1, 0xb3, 1, 2, 3, 4}},
+			},
+		},
+		[]*astits.Packet{
+			{ // PMT
+				Header: astits.PacketHeader{
+					HasPayload:                true,
+					PayloadUnitStartIndicator: true,
+					PID:                       0,
+				},
+				Payload: append([]byte{
+					0x00, 0x00, 0xb0, 0x0d, 0x00, 0x00, 0xc1, 0x00,
+					0x00, 0x00, 0x01, 0xf0, 0x00, 0x71, 0x10, 0xd8,
+					0x78,
+				}, bytes.Repeat([]byte{0xff}, 167)...),
+			},
+			{ // PAT
+				Header: astits.PacketHeader{
+					HasPayload:                true,
+					PayloadUnitStartIndicator: true,
+					PID:                       4096,
+				},
+				Payload: append([]byte{
+					0x00, 0x02, 0xb0, 0x12, 0x00, 0x01, 0xc1, 0x00,
+					0x00, 0xe1, 0x01, 0xf0, 0x00, 0x02, 0xe1, 0x01,
+					0xf0, 0x00, 0xc4, 0xf2, 0x53, 0x9c,
+				}, bytes.Repeat([]byte{0xff}, 162)...),
+			},
+			{ // PES
+				AdaptationField: &astits.PacketAdaptationField{
+					Length:                161,
+					StuffingLength:        154,
+					RandomAccessIndicator: true,
+					HasPCR:                true,
+					PCR:                   &astits.ClockReference{Base: 2691000},
+				},
+				Header: astits.PacketHeader{
+					HasAdaptationField:        true,
+					HasPayload:                true,
+					PayloadUnitStartIndicator: true,
+					PID:                       257,
+				},
+				Payload: []byte{
+					0x00, 0x00, 0x01, 0xe0, 0x00, 0x00, 0x80, 0x80,
+					0x05, 0x21, 0x00, 0xa5, 0x65, 0xc1, 0x00, 0x00,
+					0x01, 0xb3, 0x01, 0x02, 0x03, 0x04,
+				},
+			},
+		},
+	},
 	{
 		"opus",
 		&Track{
@@ -685,7 +745,15 @@ func TestReader(t *testing.T) {
 				})
 
 			case *CodecMPEG4Video:
-				r.OnDataMPEG4Video(ca.track, func(pts int64, frame []byte) error {
+				r.OnDataMPEGxVideo(ca.track, func(pts int64, frame []byte) error {
+					require.Equal(t, ca.samples[i].pts, pts)
+					require.Equal(t, ca.samples[i].data[0], frame)
+					i++
+					return nil
+				})
+
+			case *CodecMPEG1Video:
+				r.OnDataMPEGxVideo(ca.track, func(pts int64, frame []byte) error {
 					require.Equal(t, ca.samples[i].pts, pts)
 					require.Equal(t, ca.samples[i].data[0], frame)
 					i++

diff --git a/pkg/formats/mpegts/track.go b/pkg/formats/mpegts/track.go
@@ -107,6 +107,10 @@ func (t *Track) unmarshal(dem *astits.Demuxer, es *astits.PMTElementaryStream) e
 		t.Codec = &CodecMPEG4Video{}
 		return nil
 
+	case astits.StreamTypeMPEG2Video, astits.StreamTypeMPEG1Video:
+		t.Codec = &CodecMPEG1Video{}
+		return nil
+
 	case astits.StreamTypeAACAudio:
 		conf, err := findMPEG4AudioConfig(dem, es.ElementaryPID)
 		if err != nil {