Fix AVC software decoder sometimes dropping B-frames (fixes #488)

dev/demux.html

@@ -17,11 +17,45 @@
 			source: new Mediabunny.BlobSource(file),
 		});
 
+		const track = await input.getPrimaryVideoTrack();
+		if (!track) {
+			throw new Error('The synthetic file has no video track.');
+		}
+
+		let encodedPackets = 0;
+		for await (const packet of new Mediabunny.EncodedPacketSink(track).packets()) {
+			console.log(packet.timestamp)
+			encodedPackets++;
+		}
+
+		let decodedFrames = 0;
+		let lastTimestamp = null;
+		const sink = new Mediabunny.VideoSampleSink(track, {
+			hardwareAcceleration: 'prefer-software',
+		});
+		for await (const sample of sink.samples()) {
+			console.log(sample.timestamp)
+			decodedFrames++;
+			lastTimestamp = sample.timestamp;
+			sample.close();
+		}
+
+		input.dispose();
+		const reproduced = encodedPackets === 48 && decodedFrames === 47;
+		console.log({
+			decodedFrames,
+			encodedPackets,
+			lastTimestamp,
+			reproduced,
+		});
+
+		/*
 		const track = await input.getPrimaryVideoTrack();
 		const sink = new Mediabunny.EncodedPacketSink(track);
 
 		const packet = await sink.getFirstPacket();
 		console.log(packet);
+		*/
 		/*
 
 		for await (const packet of packetSink.packets()) {

packages/server/src/video-decoder.ts

@@ -82,7 +82,11 @@ export class NodeAvVideoDecoder extends CustomVideoDecoder {
 			: null;
 		codecContext.sampleAspectRatio = new NodeAv.Rational(this.pixelAspectRatio.num, this.pixelAspectRatio.den);
 
-		const ret = await codecContext.open2();
+		// Honor inferred H.264 reorder buffering even when bitstream restrictions aren't explicitly signaled
+		// See https://github.com/Vanilagy/mediabunny/issues/488
+		const options = NodeAv.Dictionary.fromObject({ strict: NodeAv.FF_COMPLIANCE_STRICT });
+
+		const ret = await codecContext.open2(codec, options);
 		NodeAv.FFmpegError.throwIfError(ret, 'Open codec context');
 
 		this.codecContext = codecContext;

shared/bitstream.ts

@@ -55,7 +55,19 @@ export class Bitstream {
 		}
 
 		this.pos = end;
-	};
+	}
+
+	copyBits(n: number, other: Bitstream) {
+		let i = 0;
+		for (i; i < n - 7; i += 8) {
+			this.writeBits(8, other.readBits(8));
+		}
+
+		const leftover = n - i;
+		if (leftover > 0) {
+			this.writeBits(leftover, other.readBits(leftover));
+		}
+	}
 
 	readAlignedByte() {
 		if (this.pos % 8 !== 0) {

src/codec-data.ts

@@ -25,6 +25,7 @@ import {
 	isChromium,
 	popcount,
 	setUint24,
+	writeExpGolomb,
 } from './misc';
 import { Logging } from './logging';
 import { PacketType } from './packet';
@@ -185,6 +186,23 @@ const removeEmulationPreventionBytes = (data: Uint8Array) => {
 	return new Uint8Array(result);
 };
 
+export const addEmulationPreventionBytes = (data: Uint8Array) => {
+	const result: number[] = [];
+	let zeroCount = 0;
+
+	for (const byte of data) {
+		if (zeroCount === 2 && byte <= 0x03) {
+			result.push(0x03);
+			zeroCount = 0;
+		}
+
+		result.push(byte);
+		zeroCount = byte === 0 ? zeroCount + 1 : 0;
+	}
+
+	return new Uint8Array(result);
+};
+
 const ANNEX_B_START_CODE = new Uint8Array([0, 0, 0, 1]);
 
 export const concatNalUnitsInAnnexB = (nalUnits: Uint8Array[]) => {
@@ -363,12 +381,14 @@ export const serializeAvcDecoderConfigurationRecord = (record: AvcDecoderConfigu
 	}
 
 	if (
-		record.avcProfileIndication === 100
-		|| record.avcProfileIndication === 110
-		|| record.avcProfileIndication === 122
-		|| record.avcProfileIndication === 144
+		(
+			record.avcProfileIndication === 100
+			|| record.avcProfileIndication === 110
+			|| record.avcProfileIndication === 122
+			|| record.avcProfileIndication === 144
+		)
+		&& record.chromaFormat !== null // Can happen if the data was too short
 	) {
-		assert(record.chromaFormat !== null);
 		assert(record.bitDepthLumaMinus8 !== null);
 		assert(record.bitDepthChromaMinus8 !== null);
 		assert(record.sequenceParameterSetExt !== null);
@@ -485,6 +505,7 @@ export const deserializeAvcDecoderConfigurationRecord = (data: Uint8Array): AvcD
 };
 
 export type AvcSpsInfo = {
+	emulationUnpreventedBytes: Uint8Array;
 	profileIdc: number;
 	constraintFlags: number;
 	levelIdc: number;
@@ -503,6 +524,9 @@ export type AvcSpsInfo = {
 	fullRangeFlag: number;
 	numReorderFrames: number;
 	maxDecFrameBuffering: number;
+	vuiParametersFlagBitOffset: number;
+	bitstreamRestrictionFlagBitOffset: number | null;
+	bitstreamRestrictionFlag: number | null;
 };
 
 const AVC_HEVC_ASPECT_RATIO_IDC_TABLE: Partial<Record<number, Rational>> = {
@@ -527,7 +551,8 @@ const AVC_HEVC_ASPECT_RATIO_IDC_TABLE: Partial<Record<number, Rational>> = {
 /** Parses an AVC SPS (Sequence Parameter Set) to extract basic information. */
 export const parseAvcSps = (sps: Uint8Array): AvcSpsInfo | null => {
 	try {
-		const bitstream = new Bitstream(removeEmulationPreventionBytes(sps));
+		const emulationUnpreventedBytes = removeEmulationPreventionBytes(sps);
+		const bitstream = new Bitstream(emulationUnpreventedBytes);
 
 		bitstream.skipBits(1); // forbidden_zero_bit
 		bitstream.skipBits(2); // nal_ref_idc
@@ -660,7 +685,10 @@ export const parseAvcSps = (sps: Uint8Array): AvcSpsInfo | null => {
 
 		let numReorderFrames: number | null = null;
 		let maxDecFrameBuffering: number | null = null;
+		let bitstreamRestrictionFlagBitOffset: number | null = null;
+		let bitstreamRestrictionFlag: number | null = null;
 
+		const vuiParametersFlagBitOffset = bitstream.pos;
 		const vuiParametersPresentFlag = bitstream.readBits(1);
 		if (vuiParametersPresentFlag) {
 			const aspectRatioInfoPresentFlag = bitstream.readBits(1);
@@ -726,7 +754,8 @@ export const parseAvcSps = (sps: Uint8Array): AvcSpsInfo | null => {
 
 			bitstream.skipBits(1); // pic_struct_present_flag
 
-			const bitstreamRestrictionFlag = bitstream.readBits(1);
+			bitstreamRestrictionFlagBitOffset = bitstream.pos;
+			bitstreamRestrictionFlag = bitstream.readBits(1);
 			if (bitstreamRestrictionFlag) {
 				bitstream.skipBits(1); // motion_vectors_over_pic_boundaries_flag
 				readExpGolomb(bitstream); // max_bytes_per_pic_denom
@@ -776,6 +805,7 @@ export const parseAvcSps = (sps: Uint8Array): AvcSpsInfo | null => {
 		assert(maxDecFrameBuffering !== null);
 
 		return {
+			emulationUnpreventedBytes,
 			profileIdc,
 			constraintFlags,
 			levelIdc,
@@ -794,6 +824,9 @@ export const parseAvcSps = (sps: Uint8Array): AvcSpsInfo | null => {
 			fullRangeFlag,
 			numReorderFrames,
 			maxDecFrameBuffering,
+			vuiParametersFlagBitOffset,
+			bitstreamRestrictionFlagBitOffset,
+			bitstreamRestrictionFlag,
 		};
 	} catch (error) {
 		Logging._error('Error parsing AVC SPS:', error);
@@ -818,6 +851,57 @@ const skipAvcHrdParameters = (bitstream: Bitstream) => {
 	bitstream.skipBits(5); // time_offset_length
 };
 
+/**
+ * Adds the missing "bitstream restriction" section within the VUI parameter section. This is done to communicate
+ * frame reorder buffer size to the decoder, which may otherwise, in its absence, assume no B-frames and drop or skip
+ * them.
+ * See https://github.com/Vanilagy/mediabunny/issues/488
+ */
+export const addAvcBitstreamRestriction = (sps: AvcSpsInfo) => {
+	assert(sps.bitstreamRestrictionFlag !== 1);
+
+	const modifiedBytes = new Uint8Array(sps.emulationUnpreventedBytes.byteLength + 64);
+	const oldBitstream = new Bitstream(sps.emulationUnpreventedBytes);
+	const newBitstream = new Bitstream(modifiedBytes);
+
+	if (sps.bitstreamRestrictionFlag === null) {
+		// No VUI at all; let's write a minimal one
+		newBitstream.copyBits(sps.vuiParametersFlagBitOffset, oldBitstream);
+		newBitstream.writeBits(1, 1); // vui_parameters_present_flag
+		newBitstream.writeBits(1, 0); // aspect_ratio_info_present_flag
+		newBitstream.writeBits(1, 0); // overscan_info_present_flag
+		newBitstream.writeBits(1, 0); // video_signal_type_present_flag
+		newBitstream.writeBits(1, 0); // chroma_loc_info_present_flag
+		newBitstream.writeBits(1, 0); // timing_info_present_flag
+		newBitstream.writeBits(1, 0); // nal_hrd_parameters_present_flag
+		newBitstream.writeBits(1, 0); // vcl_hrd_parameters_present_flag
+		newBitstream.writeBits(1, 0); // pic_struct_present_flag
+	} else {
+		// We have a VUI but no bitstream restriction info
+		assert(sps.bitstreamRestrictionFlagBitOffset !== null);
+		newBitstream.copyBits(sps.bitstreamRestrictionFlagBitOffset, oldBitstream);
+	}
+	newBitstream.writeBits(1, 1); // bitstream_restriction_flag
+
+	// Defaults from the H.264 spec:
+	newBitstream.writeBits(1, 1); // motion_vectors_over_pic_boundaries_flag
+	writeExpGolomb(newBitstream, 2); // max_bytes_per_pic_denom
+	writeExpGolomb(newBitstream, 1); // max_bits_per_mb_denom
+	writeExpGolomb(newBitstream, 16); // log2_max_mv_length_horizontal
+	writeExpGolomb(newBitstream, 16); // log2_max_mv_length_vertical
+	writeExpGolomb(newBitstream, sps.numReorderFrames);
+	writeExpGolomb(newBitstream, sps.maxDecFrameBuffering);
+
+	// There's nothing after this (VUI is at the end of SPS)
+	newBitstream.writeBits(1, 1); // rbsp_stop_one_bit
+	newBitstream.writeBits((8 - newBitstream.pos % 8) % 8, 0);
+
+	const byteLength = newBitstream.pos / 8;
+	assert(Number.isInteger(byteLength));
+
+	return addEmulationPreventionBytes(modifiedBytes.subarray(0, byteLength));
+};
+
 // Data specified in ISO 14496-15
 export type HevcDecoderConfigurationRecord = {
 	configurationVersion: number;

src/media-sink.ts

@@ -8,6 +8,7 @@
 
 import { parsePcmCodec, PCM_AUDIO_CODECS, PcmAudioCodec, VideoCodec, AudioCodec } from './codec';
 import {
+	addAvcBitstreamRestriction,
 	AvcNalUnitType,
 	concatAvcNalUnits,
 	deserializeAvcDecoderConfigurationRecord,
@@ -19,6 +20,7 @@ import {
 	iterateHevcNalUnits,
 	parseAvcSps,
 	sanitizeHevcPacketForChromium,
+	serializeAvcDecoderConfigurationRecord,
 } from './codec-data';
 import { CustomVideoDecoder, customVideoDecoders, CustomAudioDecoder, customAudioDecoders } from './custom-coder';
 import { InputDisposedError } from './input';
@@ -959,19 +961,31 @@ class VideoDecoderWrapper extends DecoderWrapper<VideoSample> {
 
 			if (isChromium()) {
 				if (codec === 'avc' && this.decoderConfig.description) {
-					// Chromium has/had a bug with playing interlaced AVC (https://issues.chromium.org/issues/456919096)
-					// which can be worked around by requesting that software decoding be used. So, here we peek into
-					// the AVC description, if present, and switch to software decoding if we find interlaced content.
 					const record = deserializeAvcDecoderConfigurationRecord(
 						toUint8Array(this.decoderConfig.description),
 					);
 					if (record && record.sequenceParameterSets.length > 0) {
 						const sps = parseAvcSps(record.sequenceParameterSets[0]!);
-						if (sps && sps.frameMbsOnlyFlag === 0) {
-							this.decoderConfig = {
-								...this.decoderConfig,
-								hardwareAcceleration: 'prefer-software',
-							};
+						if (sps) {
+							if (sps.frameMbsOnlyFlag === 0) {
+								// Chromium has/had a bug with playing interlaced AVC
+								// (https://issues.chromium.org/issues/456919096) which can be worked around by
+								// requesting that software decoding be used. So, here we peek into the AVC description,
+								// if present, and switch to software decoding if we find interlaced content.
+								this.decoderConfig = {
+									...this.decoderConfig,
+									hardwareAcceleration: 'prefer-software',
+								};
+							}
+
+							if (sps.maxDecFrameBuffering !== 0 && sps.bitstreamRestrictionFlag !== 1) {
+								// Modify the SPS to fix potential loss of B frames
+								record.sequenceParameterSets[0] = addAvcBitstreamRestriction(sps);
+								this.decoderConfig = {
+									...this.decoderConfig,
+									description: serializeAvcDecoderConfigurationRecord(record),
+								};
+							}
 						}
 					}
 				}
@@ -1077,6 +1091,23 @@ class VideoDecoderWrapper extends DecoderWrapper<VideoSample> {
 						}
 					}
 
+					if (!this.decoderConfig.description) {
+						// Do SPS fixups if necessary
+						for (let i = 0; i < filteredNalUnits.length; i++) {
+							const nalUnit = filteredNalUnits[i]!;
+							if (extractNalUnitTypeForAvc(nalUnit[0]!) !== AvcNalUnitType.SPS) {
+								continue;
+							}
+
+							const sps = parseAvcSps(nalUnit);
+							if (sps && sps.maxDecFrameBuffering !== 0 && sps.bitstreamRestrictionFlag !== 1) {
+								filteredNalUnits[i] = addAvcBitstreamRestriction(sps);
+							}
+
+							break;
+						}
+					}
+
 					const newData = concatAvcNalUnits(filteredNalUnits, this.decoderConfig);
 					packet = new EncodedPacket(newData, packet.type, packet.timestamp, packet.duration);
 				} else if (this.codec === 'hevc') {

src/misc.ts

@@ -61,6 +61,14 @@ export const readExpGolomb = (bitstream: Bitstream) => {
 	return result;
 };
 
+export const writeExpGolomb = (bitstream: Bitstream, value: number) => {
+	const codeNum = value + 1;
+	const leadingZeroBits = Math.floor(Math.log2(codeNum));
+	bitstream.writeBits(leadingZeroBits, 0);
+	bitstream.writeBits(1, 1);
+	bitstream.writeBits(leadingZeroBits, codeNum - 2 ** leadingZeroBits);
+};
+
 /** Reads a signed exponential-Golomb universal code from a Bitstream. */
 export const readSignedExpGolomb = (bitstream: Bitstream) => {
 	const codeNum = readExpGolomb(bitstream);

test/browser/media-sinks.test.ts

@@ -1,9 +1,13 @@
 import { expect, test } from 'vitest';
 import { Input } from '../../src/input.js';
-import { UrlSource } from '../../src/source.js';
+import { BufferSource, UrlSource } from '../../src/source.js';
 import { ALL_FORMATS } from '../../src/input-format.js';
 import { assert } from '../../src/misc.js';
-import { AudioSampleSink } from '../../src/media-sink.js';
+import { AudioSampleSink, VideoSampleSink } from '../../src/media-sink.js';
+import { Output } from '../../src/output.js';
+import { MpegTsOutputFormat } from '../../src/output-format.js';
+import { BufferTarget } from '../../src/target.js';
+import { Conversion } from '../../src/conversion.js';
 
 // https://github.com/Vanilagy/mediabunny/issues/370
 test('Negative audio timestamps are preserved', async () => {
@@ -24,3 +28,61 @@ test('Negative audio timestamps are preserved', async () => {
 		break;
 	}
 });
+
+test('No B-frames are skipped when software-decoding AVC', async () => {
+	using input = new Input({
+		source: new UrlSource('/missing-reorder-metadata-v1.mp4'),
+		formats: ALL_FORMATS,
+	});
+
+	const track = await input.getPrimaryVideoTrack();
+	assert(track);
+
+	const sink = new VideoSampleSink(track, {
+		hardwareAcceleration: 'prefer-software',
+	});
+	let count = 0;
+
+	// eslint-disable-next-line @typescript-eslint/no-unused-vars
+	for await (using sample of sink.samples()) {
+		count++;
+	}
+
+	expect(count).toBe(48);
+});
+
+test('No B-frames are skipped when software-decoding AVC, Annex B edition', async () => {
+	using input = new Input({
+		source: new UrlSource('/missing-reorder-metadata-v1.mp4'),
+		formats: ALL_FORMATS,
+	});
+
+	// Force Annex B by converting to MPEG-TS
+	const output = new Output({
+		format: new MpegTsOutputFormat(),
+		target: new BufferTarget(),
+	});
+
+	const conversion = await Conversion.init({ input, output, copy: { mode: 'forced' } });
+	await conversion.execute();
+
+	using newInput = new Input({
+		source: new BufferSource(output.target.buffer!),
+		formats: ALL_FORMATS,
+	});
+
+	const track = await newInput.getPrimaryVideoTrack();
+	assert(track);
+
+	const sink = new VideoSampleSink(track, {
+		hardwareAcceleration: 'prefer-software',
+	});
+	let count = 0;
+
+	// eslint-disable-next-line @typescript-eslint/no-unused-vars
+	for await (using sample of sink.samples()) {
+		count++;
+	}
+
+	expect(count).toBe(48);
+});

test/node/server-extension.test.ts

@@ -1223,6 +1223,28 @@ describe('Video', async () => {
 		await decoder.close();
 	});
 
+	test('No B-frames are skipped when decoding AVC', async () => {
+		using input = new Input({
+			source: new FilePathSource('./test/public/missing-reorder-metadata-v1.mp4'),
+			formats: ALL_FORMATS,
+		});
+
+		const track = await input.getPrimaryVideoTrack();
+		assert(track);
+
+		const sink = new VideoSampleSink(track, {
+			hardwareAcceleration: 'prefer-software',
+		});
+		let count = 0;
+
+		// eslint-disable-next-line @typescript-eslint/no-unused-vars
+		for await (using sample of sink.samples()) {
+			count++;
+		}
+
+		expect(count).toBe(48);
+	});
+
 	describe('VideoSample transformation', () => {
 		// 400x400 image: red everywhere, with a 200x200 blue square filling the bottom-left quadrant.
 		const TEST_IMAGE = (() => {