From fd2ea22e47e12ecf9719580f6e1f57a0a6080dd1 Mon Sep 17 00:00:00 2001
From: sheenachhabra <sheenachhabra@google.com>
Date: Wed, 31 Jan 2024 03:59:38 -0800
Subject: [PATCH] Allow AV1 output mime type in FrameworkMuxer

The `MediaMuxer` (FrameworkMuxer) supports AV1 mime type from API 34.
For this to work track `Format/MediaFormat` must have `CSD`
data.

Change also include providing AV1 option in demo app.

Transmuxing of an AV1 mime type input will not work because
`Extractor` does not extract AV1 CSD data.

Verified changes manually via demo app.

PiperOrigin-RevId: 603002380
---
 .../transformer/ConfigurationActivity.java    |  3 ++
 .../transformer/TransformerEndToEndTest.java  | 29 +++++++++++++++++++
 .../media3/transformer/FrameworkMuxer.java    | 23 +++++++++------
 3 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/demos/transformer/src/main/java/androidx/media3/demo/transformer/ConfigurationActivity.java b/demos/transformer/src/main/java/androidx/media3/demo/transformer/ConfigurationActivity.java
index 433704cbd7..ae0d1923f9 100644
--- a/demos/transformer/src/main/java/androidx/media3/demo/transformer/ConfigurationActivity.java
+++ b/demos/transformer/src/main/java/androidx/media3/demo/transformer/ConfigurationActivity.java
@@ -316,6 +316,9 @@ public final class ConfigurationActivity extends AppCompatActivity {
     if (SDK_INT >= 24) {
       videoMimeAdapter.add(MimeTypes.VIDEO_H265);
     }
+    if (SDK_INT >= 34) {
+      videoMimeAdapter.add(MimeTypes.VIDEO_AV1);
+    }
 
     ArrayAdapter<String> resolutionHeightAdapter =
         new ArrayAdapter<>(/* context= */ this, R.layout.spinner_item);
diff --git a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerEndToEndTest.java b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerEndToEndTest.java
index 971efef128..778aacecd9 100644
--- a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerEndToEndTest.java
+++ b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/TransformerEndToEndTest.java
@@ -49,6 +49,7 @@ import androidx.media3.common.C;
 import androidx.media3.common.Effect;
 import androidx.media3.common.Format;
 import androidx.media3.common.MediaItem;
+import androidx.media3.common.MimeTypes;
 import androidx.media3.common.OnInputFrameProcessedListener;
 import androidx.media3.common.VideoFrameProcessingException;
 import androidx.media3.common.audio.AudioProcessor;
@@ -923,6 +924,34 @@ public class TransformerEndToEndTest {
     assertThat(result.exportResult.channelCount).isEqualTo(2);
   }
 
+  @Test
+  public void transcode_withOutputMimeTypeAv1_completesSuccessfully() throws Exception {
+    String testId = "transcode_withOutputMimeTypeAv1_completesSuccessfully";
+    if (AndroidTestUtil.skipAndLogIfFormatsUnsupported(
+        context,
+        testId,
+        /* inputFormat= */ MP4_ASSET_FORMAT,
+        /* outputFormat= */ MP4_ASSET_FORMAT
+            .buildUpon()
+            .setSampleMimeType(MimeTypes.VIDEO_AV1)
+            .setCodecs(null)
+            .build())) {
+      return;
+    }
+    MediaItem mediaItem = MediaItem.fromUri(Uri.parse(MP4_ASSET_URI_STRING));
+    EditedMediaItem editedMediaItem = new EditedMediaItem.Builder(mediaItem).build();
+    Transformer transformer =
+        new Transformer.Builder(context).setVideoMimeType(MimeTypes.VIDEO_AV1).build();
+
+    ExportTestResult exportTestResult =
+        new TransformerAndroidTestRunner.Builder(context, transformer)
+            .build()
+            .run(testId, editedMediaItem);
+
+    // TODO: b/322813915 - Assert on output mime type.
+    assertThat(exportTestResult.exportResult.exportException).isNull();
+  }
+
   private static AudioProcessor createSonic(float pitch) {
     SonicAudioProcessor sonic = new SonicAudioProcessor();
     sonic.setPitch(pitch);
diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/FrameworkMuxer.java b/libraries/transformer/src/main/java/androidx/media3/transformer/FrameworkMuxer.java
index 638dbc74e9..671091c382 100644
--- a/libraries/transformer/src/main/java/androidx/media3/transformer/FrameworkMuxer.java
+++ b/libraries/transformer/src/main/java/androidx/media3/transformer/FrameworkMuxer.java
@@ -39,17 +39,9 @@ import java.nio.ByteBuffer;
 
 /** {@link Muxer} implementation that uses a {@link MediaMuxer}. */
 /* package */ final class FrameworkMuxer implements Muxer {
-
   // MediaMuxer supported sample formats are documented in MediaMuxer.addTrack(MediaFormat).
   private static final ImmutableList<String> SUPPORTED_VIDEO_SAMPLE_MIME_TYPES =
-      Util.SDK_INT >= 24
-          ? ImmutableList.of(
-              MimeTypes.VIDEO_H265,
-              MimeTypes.VIDEO_H264,
-              MimeTypes.VIDEO_H263,
-              MimeTypes.VIDEO_MP4V)
-          : ImmutableList.of(MimeTypes.VIDEO_H264, MimeTypes.VIDEO_H263, MimeTypes.VIDEO_MP4V);
-
+      getSupportedVideoSampleMimeTypes();
   private static final ImmutableList<String> SUPPORTED_AUDIO_SAMPLE_MIME_TYPES =
       ImmutableList.of(MimeTypes.AUDIO_AAC, MimeTypes.AUDIO_AMR_NB, MimeTypes.AUDIO_AMR_WB);
 
@@ -259,4 +251,17 @@ import java.nio.ByteBuffer;
       throw e;
     }
   }
+
+  private static ImmutableList<String> getSupportedVideoSampleMimeTypes() {
+    ImmutableList.Builder<String> supportedMimeTypes =
+        new ImmutableList.Builder<String>()
+            .add(MimeTypes.VIDEO_H264, MimeTypes.VIDEO_H263, MimeTypes.VIDEO_MP4V);
+    if (SDK_INT >= 24) {
+      supportedMimeTypes.add(MimeTypes.VIDEO_H265);
+    }
+    if (SDK_INT >= 34) {
+      supportedMimeTypes.add(MimeTypes.VIDEO_AV1);
+    }
+    return supportedMimeTypes.build();
+  }
 }