Process all samples before writing them to mdat

This refactoring is required to support 3 byte NAL start code, in which
the sample ByteBuffer might change after AnnexB to Avcc conversion.
This would required changes in the corresponding BufferInfo as well.

PiperOrigin-RevId: 617538338
This commit is contained in:
sheenachhabra 2024-03-20 09:29:40 -07:00 committed by Copybara-Service
parent 276e0655f4
commit 2491dd07e1

View file

@ -17,6 +17,7 @@ package androidx.media3.muxer;
import static androidx.media3.common.util.Assertions.checkArgument; import static androidx.media3.common.util.Assertions.checkArgument;
import static androidx.media3.common.util.Assertions.checkNotNull; import static androidx.media3.common.util.Assertions.checkNotNull;
import static androidx.media3.common.util.Assertions.checkState;
import static androidx.media3.muxer.AnnexBUtils.doesSampleContainAnnexBNalUnits; import static androidx.media3.muxer.AnnexBUtils.doesSampleContainAnnexBNalUnits;
import static androidx.media3.muxer.Boxes.BOX_HEADER_SIZE; import static androidx.media3.muxer.Boxes.BOX_HEADER_SIZE;
import static androidx.media3.muxer.Boxes.MFHD_BOX_CONTENT_SIZE; import static androidx.media3.muxer.Boxes.MFHD_BOX_CONTENT_SIZE;
@ -226,12 +227,12 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
} }
output.write(Boxes.moof(Boxes.mfhd(currentFragmentSequenceNumber), trafBoxes)); output.write(Boxes.moof(Boxes.mfhd(currentFragmentSequenceNumber), trafBoxes));
writeMdatBox(); writeMdatBox(trackInfos);
currentFragmentSequenceNumber++; currentFragmentSequenceNumber++;
} }
private void writeMdatBox() throws IOException { private void writeMdatBox(List<ProcessedTrackInfo> trackInfos) throws IOException {
long mdatStartPosition = output.position(); long mdatStartPosition = output.position();
int mdatHeaderSize = 8; // 4 bytes (box size) + 4 bytes (box name) int mdatHeaderSize = 8; // 4 bytes (box size) + 4 bytes (box name)
ByteBuffer header = ByteBuffer.allocate(mdatHeaderSize); ByteBuffer header = ByteBuffer.allocate(mdatHeaderSize);
@ -241,17 +242,12 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
output.write(header); output.write(header);
long bytesWritten = 0; long bytesWritten = 0;
for (int i = 0; i < tracks.size(); i++) { for (int trackInfoIndex = 0; trackInfoIndex < trackInfos.size(); trackInfoIndex++) {
Track currentTrack = tracks.get(i); ProcessedTrackInfo currentTrackInfo = trackInfos.get(trackInfoIndex);
while (!currentTrack.pendingSamplesByteBuffer.isEmpty()) { for (int sampleIndex = 0;
ByteBuffer currentSampleByteBuffer = currentTrack.pendingSamplesByteBuffer.removeFirst(); sampleIndex < currentTrackInfo.pendingSamplesByteBuffer.size();
sampleIndex++) {
// Convert the H.264/H.265 samples from Annex-B format (output by MediaCodec) to bytesWritten += output.write(currentTrackInfo.pendingSamplesByteBuffer.get(sampleIndex));
// Avcc format (required by MP4 container).
if (doesSampleContainAnnexBNalUnits(checkNotNull(currentTrack.format.sampleMimeType))) {
annexBToAvccConverter.process(currentSampleByteBuffer);
}
bytesWritten += output.write(currentSampleByteBuffer);
} }
} }
@ -280,42 +276,66 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
} }
private ProcessedTrackInfo processTrack(int trackId, Track track) { private ProcessedTrackInfo processTrack(int trackId, Track track) {
List<BufferInfo> sampleBufferInfos = new ArrayList<>(track.pendingSamplesBufferInfo); checkState(track.pendingSamplesByteBuffer.size() == track.pendingSamplesBufferInfo.size());
ImmutableList.Builder<ByteBuffer> pendingSamplesByteBuffer = new ImmutableList.Builder<>();
if (doesSampleContainAnnexBNalUnits(checkNotNull(track.format.sampleMimeType))) {
while (!track.pendingSamplesByteBuffer.isEmpty()) {
ByteBuffer currentSampleByteBuffer = track.pendingSamplesByteBuffer.removeFirst();
annexBToAvccConverter.process(currentSampleByteBuffer);
pendingSamplesByteBuffer.add(currentSampleByteBuffer);
}
} else {
pendingSamplesByteBuffer.addAll(track.pendingSamplesByteBuffer);
track.pendingSamplesByteBuffer.clear();
}
ImmutableList.Builder<BufferInfo> pendingSamplesBufferInfoBuilder =
new ImmutableList.Builder<>();
pendingSamplesBufferInfoBuilder.addAll(track.pendingSamplesBufferInfo);
track.pendingSamplesBufferInfo.clear();
ImmutableList<BufferInfo> pendingSamplesBufferInfo = pendingSamplesBufferInfoBuilder.build();
List<Long> sampleDurations = List<Long> sampleDurations =
Boxes.convertPresentationTimestampsToDurationsVu( Boxes.convertPresentationTimestampsToDurationsVu(
sampleBufferInfos, pendingSamplesBufferInfo,
/* firstSamplePresentationTimeUs= */ currentFragmentSequenceNumber == 1 /* firstSamplePresentationTimeUs= */ currentFragmentSequenceNumber == 1
? minInputPresentationTimeUs ? minInputPresentationTimeUs
: sampleBufferInfos.get(0).presentationTimeUs, : pendingSamplesBufferInfo.get(0).presentationTimeUs,
track.videoUnitTimebase(), track.videoUnitTimebase(),
Mp4Muxer.LAST_FRAME_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION); Mp4Muxer.LAST_FRAME_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION);
ImmutableList.Builder<SampleMetadata> pendingSamplesMetadata = new ImmutableList.Builder<>(); ImmutableList.Builder<SampleMetadata> pendingSamplesMetadata = new ImmutableList.Builder<>();
int totalSamplesSize = 0; int totalSamplesSize = 0;
for (int i = 0; i < sampleBufferInfos.size(); i++) { for (int i = 0; i < pendingSamplesBufferInfo.size(); i++) {
totalSamplesSize += sampleBufferInfos.get(i).size; totalSamplesSize += pendingSamplesBufferInfo.get(i).size;
pendingSamplesMetadata.add( pendingSamplesMetadata.add(
new SampleMetadata( new SampleMetadata(
sampleDurations.get(i), sampleDurations.get(i),
sampleBufferInfos.get(i).size, pendingSamplesBufferInfo.get(i).size,
sampleBufferInfos.get(i).flags)); pendingSamplesBufferInfo.get(i).flags));
} }
// Clear the queue. return new ProcessedTrackInfo(
track.pendingSamplesBufferInfo.clear(); trackId,
return new ProcessedTrackInfo(trackId, totalSamplesSize, pendingSamplesMetadata.build()); totalSamplesSize,
pendingSamplesByteBuffer.build(),
pendingSamplesMetadata.build());
} }
private static class ProcessedTrackInfo { private static class ProcessedTrackInfo {
public final int trackId; public final int trackId;
public final int totalSamplesSize; public final int totalSamplesSize;
public final ImmutableList<ByteBuffer> pendingSamplesByteBuffer;
public final ImmutableList<SampleMetadata> pendingSamplesMetadata; public final ImmutableList<SampleMetadata> pendingSamplesMetadata;
public ProcessedTrackInfo( public ProcessedTrackInfo(
int trackId, int totalSamplesSize, ImmutableList<SampleMetadata> pendingSamplesMetadata) { int trackId,
int totalSamplesSize,
ImmutableList<ByteBuffer> pendingSamplesByteBuffer,
ImmutableList<SampleMetadata> pendingSamplesMetadata) {
this.trackId = trackId; this.trackId = trackId;
this.totalSamplesSize = totalSamplesSize; this.totalSamplesSize = totalSamplesSize;
this.pendingSamplesByteBuffer = pendingSamplesByteBuffer;
this.pendingSamplesMetadata = pendingSamplesMetadata; this.pendingSamplesMetadata = pendingSamplesMetadata;
} }
} }