fix(avc2mp4): produce playable mp4 when timestamps are out of order

pkg/avc2mp4/mp4writer.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 	"io"
 	"math"
+	"slices"
 	"time"
 
 	"github.com/yapingcat/gomedia/go-mp4"
@@ -18,6 +19,7 @@ type ConvertResult struct {
 type accessUnit struct {
 	nalus       []NALUnit
 	timestampUs uint64
+	hasSlice    bool
 }
 
 const (
@@ -39,8 +41,11 @@ func Convert(avcData []byte, output io.WriteSeeker) (*ConvertResult, error) {
 		return nil, err
 	}
 
-	firstTs := units[0].timestampUs
-	lastTs := units[len(units)-1].timestampUs
+	// units are in decode order, which with B-frames is not presentation order
+	firstTs, lastTs := units[0].timestampUs, units[0].timestampUs
+	for _, au := range units {
+		firstTs, lastTs = min(firstTs, au.timestampUs), max(lastTs, au.timestampUs)
+	}
 	var duration time.Duration
 	const maxDurationMicros = uint64(math.MaxInt64) / uint64(time.Microsecond)
 	if delta := lastTs - firstTs; lastTs > firstTs && delta <= maxDurationMicros {
@@ -53,11 +58,22 @@ func Convert(avcData []byte, output io.WriteSeeker) (*ConvertResult, error) {
 	}, nil
 }
 
+// groupAccessUnits splits the stream into one access unit per picture. the
+// encoder emits our timecode SEI *before* the picture's slice, so a timestamp is
+// held until the next slice claims it. everything before the first SPS is
+// undecodable and dropped, as are pictures that never got a timestamp.
 func groupAccessUnits(nalus []NALUnit) []accessUnit {
 	var units []accessUnit
 	var current *accessUnit
+	var pendingTs uint64
 	seenSPS := false
 
+	flush := func() {
+		if current != nil && current.hasSlice && current.timestampUs > 0 {
+			units = append(units, *current)
+		}
+	}
+
 	for _, nalu := range nalus {
 		if nalu.Type == nalTypeSPS {
 			seenSPS = true
@@ -66,38 +82,59 @@ func groupAccessUnits(nalus []NALUnit) []accessUnit {
 			continue
 		}
 
-		// check if this NAL starts a new access unit
-		isSlice := nalu.Type == 1 || nalu.Type == 5
-		isSPS := nalu.Type == nalTypeSPS
-
-		if isSlice || isSPS {
-			if current != nil && current.timestampUs > 0 {
-				units = append(units, *current)
+		if nalu.Type == nalTypeSEI {
+			if ts, ok := ParseTimestamp(nalu.Data); ok {
+				pendingTs = ts
+				continue // don't include custom SEI in muxed output
 			}
-			current = &accessUnit{}
 		}
 
-		if current == nil {
+		isSlice := nalu.Type == 1 || nalu.Type == 5
+		// a slice with no fresh timestamp is another slice of the same picture
+		isSamePicture := isSlice && pendingTs == 0
+		if current == nil || (current.hasSlice && !isSamePicture) {
+			flush()
 			current = &accessUnit{}
 		}
 
-		// extract timestamp from our custom SEI
-		if nalu.Type == nalTypeSEI {
-			if ts, ok := ParseTimestamp(nalu.Data); ok {
-				current.timestampUs = ts
-				continue // don't include custom SEI in muxed output
-			}
+		current.nalus = append(current.nalus, nalu)
+		if isSlice && !current.hasSlice {
+			current.hasSlice = true
+			current.timestampUs = pendingTs
+			pendingTs = 0
 		}
+	}
 
-		current.nalus = append(current.nalus, nalu)
+	flush()
+	return units
+}
+
+// sampleTimesMs returns pts and dts in milliseconds for units given in decode
+// order. with B-frames the encoder emits pictures out of presentation order, so
+// dts cannot simply equal pts: dts walks the sorted timestamps (monotonic), and
+// every pts is pushed back by the largest reorder delay so dts <= pts holds.
+// without reordering the delay is zero and dts == pts.
+func sampleTimesMs(units []accessUnit) (pts []uint64, dts []uint64) {
+	dts = make([]uint64, len(units))
+	for i, au := range units {
+		dts[i] = au.timestampUs
 	}
+	slices.Sort(dts)
 
-	// flush last access unit
-	if current != nil && current.timestampUs > 0 {
-		units = append(units, *current)
+	first := dts[0]
+	var delayUs uint64
+	for i, au := range units {
+		if dts[i] > au.timestampUs {
+			delayUs = max(delayUs, dts[i]-au.timestampUs)
+		}
 	}
 
-	return units
+	pts = make([]uint64, len(units))
+	for i, au := range units {
+		pts[i] = (au.timestampUs - first + delayUs) / 1000
+		dts[i] = (dts[i] - first) / 1000
+	}
+	return pts, dts
 }
 
 func writeMp4(units []accessUnit, output io.WriteSeeker) error {
@@ -108,14 +145,9 @@ func writeMp4(units []accessUnit, output io.WriteSeeker) error {
 
 	trackID := muxer.AddVideoTrack(mp4.MP4_CODEC_H264)
 
-	firstTs := units[0].timestampUs
-
-	for _, au := range units {
-		annexB := buildAnnexB(au.nalus)
-		ptsMs := (au.timestampUs - firstTs) / 1000
-		dtsMs := ptsMs // baseline profile, no B-frames
-
-		err := muxer.Write(trackID, annexB, ptsMs, dtsMs)
+	pts, dts := sampleTimesMs(units)
+	for i, au := range units {
+		err := muxer.Write(trackID, buildAnnexB(au.nalus), pts[i], dts[i])
 		if err != nil {
 			return fmt.Errorf("writing frame: %w", err)
 		}

pkg/avc2mp4/mp4writer_test.go

@@ -2,6 +2,7 @@ package avc2mp4
 
 import (
 	"bytes"
+	"slices"
 	"testing"
 )
 
@@ -21,44 +22,66 @@ func nalu(nalType byte) NALUnit {
 	return NALUnit{Type: nalType, Data: []byte{nalType, 0xAA}}
 }
 
-// grouping closes an access unit when the next slice or SPS arrives, so a unit
-// is a run of NAL units followed by the SEI carrying its timestamp
+// the encoder emits the timecode SEI right before the picture it stamps, so a
+// timestamp belongs to the slice that follows it
+func timestampsOf(units []accessUnit) []uint64 {
+	timestamps := make([]uint64, len(units))
+	for i, unit := range units {
+		timestamps[i] = unit.timestampUs
+	}
+	return timestamps
+}
+
 func TestGroupAccessUnitsIgnoresNalusBeforeFirstSPS(t *testing.T) {
 	// a recording can start mid-stream; everything before the first SPS is
 	// undecodable and must be dropped
 	units := groupAccessUnits([]NALUnit{
-		nalu(1),
 		seiWithTimestamp(1_000_000),
+		nalu(1),
 		nalu(nalTypeSPS),
-		nalu(5),
 		seiWithTimestamp(2_000_000),
+		nalu(5),
 	})
 
-	if len(units) != 1 {
-		t.Fatalf("expected 1 access unit, got %d", len(units))
-	}
-	if units[0].timestampUs != 2_000_000 {
-		t.Errorf("expected timestamp 2000000, got %d", units[0].timestampUs)
+	if got := timestampsOf(units); !slices.Equal(got, []uint64{2_000_000}) {
+		t.Fatalf("expected only the keyframe at 2000000, got %v", got)
 	}
 }
 
-func TestGroupAccessUnitsSplitsOnEachSlice(t *testing.T) {
+func TestGroupAccessUnitsGivesEachSliceTheTimestampBeforeIt(t *testing.T) {
 	units := groupAccessUnits([]NALUnit{
 		nalu(nalTypeSPS),
-		nalu(5),
 		seiWithTimestamp(1_000_000),
-		nalu(1),
+		nalu(5),
 		seiWithTimestamp(2_000_000),
 		nalu(1),
 		seiWithTimestamp(3_000_000),
+		nalu(1),
+	})
+
+	if got := timestampsOf(units); !slices.Equal(got, []uint64{1_000_000, 2_000_000, 3_000_000}) {
+		t.Fatalf("expected one access unit per slice, got %v", got)
+	}
+}
+
+func TestGroupAccessUnitsKeepsParameterSetsWithTheirKeyframe(t *testing.T) {
+	units := groupAccessUnits([]NALUnit{
+		nalu(nalTypeSPS),
+		nalu(8),
+		seiWithTimestamp(1_000_000),
+		nalu(5),
+		nalu(nalTypeSPS),
+		nalu(8),
+		seiWithTimestamp(2_000_000),
+		nalu(5),
 	})
 
-	if len(units) != 3 {
-		t.Fatalf("expected 3 access units, got %d", len(units))
+	if len(units) != 2 {
+		t.Fatalf("expected 2 access units, got %d", len(units))
 	}
-	for i, want := range []uint64{1_000_000, 2_000_000, 3_000_000} {
-		if units[i].timestampUs != want {
-			t.Errorf("access unit %d: expected timestamp %d, got %d", i, want, units[i].timestampUs)
+	for i, unit := range units {
+		if len(unit.nalus) != 3 || unit.nalus[0].Type != nalTypeSPS || unit.nalus[2].Type != 5 {
+			t.Errorf("access unit %d: expected SPS, PPS, IDR together, got %d nalus", i, len(unit.nalus))
 		}
 	}
 }
@@ -75,15 +98,49 @@ func TestGroupAccessUnitsDropsUnitsWithoutTimestamp(t *testing.T) {
 	}
 }
 
+// High profile encoders emit B-frames: decode order P B B B while presentation
+// order is B B B P. these are real timestamps from a devicekit-ios stream.
+func unitsAt(timestampsUs ...uint64) []accessUnit {
+	units := make([]accessUnit, len(timestampsUs))
+	for i, ts := range timestampsUs {
+		units[i] = accessUnit{timestampUs: ts}
+	}
+	return units
+}
+
+func TestSampleTimesStayMonotonicWhenFramesAreReordered(t *testing.T) {
+	pts, dts := sampleTimesMs(unitsAt(8700544384, 8700511052, 8700494386, 8700527718, 8700611048))
+
+	if !slices.IsSorted(dts) {
+		t.Errorf("dts must never go backwards, got %v", dts)
+	}
+	for i := range pts {
+		if dts[i] > pts[i] {
+			t.Errorf("sample %d: dts %d is after pts %d", i, dts[i], pts[i])
+		}
+	}
+	if longest := slices.Max(pts); longest > 1000 {
+		t.Errorf("five frames at 60fps should span well under a second, got pts up to %dms", longest)
+	}
+}
+
+func TestSampleTimesMatchWhenFramesAreInOrder(t *testing.T) {
+	pts, dts := sampleTimesMs(unitsAt(5_000_000, 5_016_000, 5_033_000))
+
+	if !slices.Equal(pts, dts) || !slices.Equal(pts, []uint64{0, 16, 33}) {
+		t.Errorf("expected pts == dts == [0 16 33], got pts=%v dts=%v", pts, dts)
+	}
+}
+
 // the timecode SEI is ours, not part of the encoded stream, so it must not be
 // muxed into the output
 func TestGroupAccessUnitsExcludesTimecodeSEIFromPayload(t *testing.T) {
 	units := groupAccessUnits([]NALUnit{
 		nalu(nalTypeSPS),
-		nalu(5),
 		seiWithTimestamp(1_000_000),
-		nalu(1),
+		nalu(5),
 		seiWithTimestamp(2_000_000),
+		nalu(1),
 	})
 
 	if len(units) == 0 {
@@ -103,11 +160,11 @@ func TestGroupAccessUnitsKeepsForeignSEIInPayload(t *testing.T) {
 	foreignSEI := NALUnit{Type: nalTypeSEI, Data: []byte{0x06, 0x01, 0x02, 0x80}}
 	units := groupAccessUnits([]NALUnit{
 		nalu(nalTypeSPS),
-		nalu(5),
-		foreignSEI,
 		seiWithTimestamp(1_000_000),
-		nalu(1),
+		foreignSEI,
+		nalu(5),
 		seiWithTimestamp(2_000_000),
+		nalu(1),
 	})
 
 	found := false