diff --git a/pkg/avc2mp4/mp4writer.go b/pkg/avc2mp4/mp4writer.go index cf5125f..d114511 100644 --- a/pkg/avc2mp4/mp4writer.go +++ b/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) } diff --git a/pkg/avc2mp4/mp4writer_test.go b/pkg/avc2mp4/mp4writer_test.go index 43d5792..f0ff189 100644 --- a/pkg/avc2mp4/mp4writer_test.go +++ b/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