From 3279f719eeb3b9d29713e07b27937f1e7e829957 Mon Sep 17 00:00:00 2001 From: gmegidish Date: Fri, 18 Sep 2026 17:49:46 +0200 Subject: [PATCH] fix(avc2mp4): produce playable mp4 when timestamps are out of order Recordings from the current devicekit-ios extension decoded as one frame with a duration of hundreds of hours. - the timecode SEI precedes the picture it stamps, but grouping attached it to the previous access unit, shifting every timestamp by one frame and emitting SPS/PPS as a sample without a slice - High profile streams contain B-frames, so timestamps arrive in decode order; writing pts == dts made the unsigned deltas wrap Timestamps are now held until the next slice claims them, parameter sets stay with their key frame, and dts walks the sorted timestamps with pts pushed back by the reorder delay so dts <= pts always holds. --- pkg/avc2mp4/mp4writer.go | 92 +++++++++++++++++++++---------- pkg/avc2mp4/mp4writer_test.go | 101 ++++++++++++++++++++++++++-------- 2 files changed, 141 insertions(+), 52 deletions(-) diff --git a/pkg/avc2mp4/mp4writer.go b/pkg/avc2mp4/mp4writer.go index cf5125f1..d1145119 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 43d57921..f0ff189d 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