From 3272ad50f3a377eaae37d4d94cdfdd8fe0d4cb97 Mon Sep 17 00:00:00 2001
From: sheenachhabra <sheenachhabra@google.com>
Date: Tue, 19 Mar 2024 09:15:44 -0700
Subject: [PATCH] Skip AnnexB to Avcc conversion for AV1 video format

The AV1 video stream does not contain NAL unit which is a concept
specific to H.264/H.265.
This was not caught before muxer does an in place replacement of
`NAL start code` with `NAL length`. In the absence of `NAL start code`
it was automatically a no-op.

PiperOrigin-RevId: 617193317
---
 .../main/java/androidx/media3/muxer/AnnexBUtils.java   | 10 ++++++++++
 .../java/androidx/media3/muxer/BasicMp4Writer.java     |  4 ++--
 .../androidx/media3/muxer/FragmentedMp4Writer.java     |  3 ++-
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/libraries/muxer/src/main/java/androidx/media3/muxer/AnnexBUtils.java b/libraries/muxer/src/main/java/androidx/media3/muxer/AnnexBUtils.java
index 7176b8896a..1addc17732 100644
--- a/libraries/muxer/src/main/java/androidx/media3/muxer/AnnexBUtils.java
+++ b/libraries/muxer/src/main/java/androidx/media3/muxer/AnnexBUtils.java
@@ -15,6 +15,7 @@
  */
 package androidx.media3.muxer;
 
+import androidx.media3.common.MimeTypes;
 import com.google.common.collect.ImmutableList;
 import java.nio.ByteBuffer;
 
@@ -87,6 +88,15 @@ import java.nio.ByteBuffer;
     return output;
   }
 
+  /**
+   * Returns whether the sample of the given MIME type will contain NAL units in Annex-B format
+   * (ISO/IEC 14496-10 Annex B, which uses start codes to delineate NAL units).
+   */
+  public static boolean doesSampleContainAnnexBNalUnits(String sampleMimeType) {
+    return sampleMimeType.equals(MimeTypes.VIDEO_H264)
+        || sampleMimeType.equals(MimeTypes.VIDEO_H265);
+  }
+
   private static ByteBuffer getBytes(ByteBuffer buf, int offset, int length) {
     ByteBuffer result = buf.duplicate();
     result.position(offset);
diff --git a/libraries/muxer/src/main/java/androidx/media3/muxer/BasicMp4Writer.java b/libraries/muxer/src/main/java/androidx/media3/muxer/BasicMp4Writer.java
index 4af85d4225..5ef82eb3ec 100644
--- a/libraries/muxer/src/main/java/androidx/media3/muxer/BasicMp4Writer.java
+++ b/libraries/muxer/src/main/java/androidx/media3/muxer/BasicMp4Writer.java
@@ -18,12 +18,12 @@ package androidx.media3.muxer;
 import static androidx.media3.common.util.Assertions.checkArgument;
 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 java.lang.Math.max;
 import static java.lang.Math.min;
 
 import android.media.MediaCodec.BufferInfo;
 import androidx.media3.common.Format;
-import androidx.media3.common.MimeTypes;
 import androidx.media3.common.util.Util;
 import androidx.media3.muxer.Mp4Muxer.TrackToken;
 import com.google.common.collect.Range;
@@ -297,7 +297,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
 
       // Convert the H.264/H.265 samples from Annex-B format (output by MediaCodec) to
       // Avcc format (required by MP4 container).
-      if (MimeTypes.isVideo(track.format.sampleMimeType)) {
+      if (doesSampleContainAnnexBNalUnits(checkNotNull(track.format.sampleMimeType))) {
         annexBToAvccConverter.process(currentSampleByteBuffer);
       }
 
diff --git a/libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java b/libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java
index dd2756690b..0948585aff 100644
--- a/libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java
+++ b/libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java
@@ -17,6 +17,7 @@ package androidx.media3.muxer;
 
 import static androidx.media3.common.util.Assertions.checkArgument;
 import static androidx.media3.common.util.Assertions.checkNotNull;
+import static androidx.media3.muxer.AnnexBUtils.doesSampleContainAnnexBNalUnits;
 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.TFHD_BOX_CONTENT_SIZE;
@@ -247,7 +248,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
         // Convert the H.264/H.265 samples from Annex-B format (output by MediaCodec) to
         // Avcc format (required by MP4 container).
-        if (MimeTypes.isVideo(currentTrack.format.sampleMimeType)) {
+        if (doesSampleContainAnnexBNalUnits(checkNotNull(currentTrack.format.sampleMimeType))) {
           annexBToAvccConverter.process(currentSampleByteBuffer);
         }
         bytesWritten += output.write(currentSampleByteBuffer);