From 0936b549ae822e9d140837163af474bc4cb92c5f Mon Sep 17 00:00:00 2001
From: dancho <dancho@google.com>
Date: Wed, 15 Jan 2025 07:24:48 -0800
Subject: [PATCH] Enable non-reference frame skipping in FrameExtractor

Speed up FrameExtractor for videos with non-ref frames

PiperOrigin-RevId: 715789816
---
 .../transformer/FrameExtractorTest.java       | 24 +++++++++++++++++++
 .../ExperimentalFrameExtractor.java           | 10 +++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/FrameExtractorTest.java b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/FrameExtractorTest.java
index 5f9b6c5b0a..3560d02863 100644
--- a/libraries/transformer/src/androidTest/java/androidx/media3/transformer/FrameExtractorTest.java
+++ b/libraries/transformer/src/androidTest/java/androidx/media3/transformer/FrameExtractorTest.java
@@ -21,7 +21,9 @@ import static androidx.media3.exoplayer.SeekParameters.CLOSEST_SYNC;
 import static androidx.media3.test.utils.BitmapPixelTestUtil.maybeSaveTestBitmap;
 import static androidx.media3.test.utils.BitmapPixelTestUtil.readBitmap;
 import static androidx.media3.test.utils.TestUtil.assertBitmapsAreSimilar;
+import static androidx.media3.transformer.AndroidTestUtil.MP4_ASSET;
 import static androidx.media3.transformer.AndroidTestUtil.MP4_TRIM_OPTIMIZATION_270;
+import static androidx.media3.transformer.AndroidTestUtil.assumeFormatsSupported;
 import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
 import static java.util.concurrent.TimeUnit.SECONDS;
@@ -473,4 +475,26 @@ public class FrameExtractorTest {
     assertThat(frameSecondItem.presentationTimeMs).isEqualTo(8_531);
     assertBitmapsAreSimilar(expectedBitmapSecondItem, actualBitmapSecondItem, PSNR_THRESHOLD);
   }
+
+  @Test
+  public void extractFrame_oneFrame_decodesReferenceFramesOnly() throws Exception {
+    assumeFormatsSupported(
+        context, testId, /* inputFormat= */ MP4_ASSET.videoFormat, /* outputFormat= */ null);
+    frameExtractor =
+        new ExperimentalFrameExtractor(
+            context, new ExperimentalFrameExtractor.Configuration.Builder().build());
+    frameExtractor.setMediaItem(
+        MediaItem.fromUri(MP4_ASSET.uri), /* effects= */ ImmutableList.of());
+
+    ListenableFuture<Frame> frameFuture = frameExtractor.getFrame(/* positionMs= */ 967);
+    Frame frame = frameFuture.get(TIMEOUT_SECONDS, SECONDS);
+
+    assertThat(frame.presentationTimeMs).isEqualTo(967);
+    assertThat(
+            frameExtractor
+                .getDecoderCounters()
+                .get(TIMEOUT_SECONDS, SECONDS)
+                .skippedInputBufferCount)
+        .isEqualTo(13);
+  }
 }
diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/ExperimentalFrameExtractor.java b/libraries/transformer/src/main/java/androidx/media3/transformer/ExperimentalFrameExtractor.java
index 5c20bd85bd..85dd8dd12e 100644
--- a/libraries/transformer/src/main/java/androidx/media3/transformer/ExperimentalFrameExtractor.java
+++ b/libraries/transformer/src/main/java/androidx/media3/transformer/ExperimentalFrameExtractor.java
@@ -49,6 +49,7 @@ import androidx.annotation.Nullable;
 import androidx.annotation.RequiresApi;
 import androidx.annotation.VisibleForTesting;
 import androidx.concurrent.futures.CallbackToFutureAdapter;
+import androidx.media3.common.C;
 import androidx.media3.common.Effect;
 import androidx.media3.common.Format;
 import androidx.media3.common.GlObjectsProvider;
@@ -80,9 +81,11 @@ import androidx.media3.exoplayer.analytics.AnalyticsListener;
 import androidx.media3.exoplayer.mediacodec.MediaCodecAdapter;
 import androidx.media3.exoplayer.mediacodec.MediaCodecInfo;
 import androidx.media3.exoplayer.mediacodec.MediaCodecSelector;
+import androidx.media3.exoplayer.source.DefaultMediaSourceFactory;
 import androidx.media3.exoplayer.source.MediaSource;
 import androidx.media3.exoplayer.video.MediaCodecVideoRenderer;
 import androidx.media3.exoplayer.video.VideoRendererEventListener;
+import androidx.media3.extractor.DefaultExtractorsFactory;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
@@ -273,6 +276,10 @@ public final class ExperimentalFrameExtractor {
    * @param configuration The {@link Configuration} for this frame extractor.
    */
   public ExperimentalFrameExtractor(Context context, Configuration configuration) {
+    MediaSource.Factory mediaSourceFactory =
+        new DefaultMediaSourceFactory(context, new DefaultExtractorsFactory())
+            .experimentalSetCodecsToParseWithinGopSampleDependencies(
+                C.VIDEO_CODEC_FLAG_H264 | C.VIDEO_CODEC_FLAG_H265);
     player =
         new ExoPlayer.Builder(
                 context,
@@ -287,7 +294,8 @@ public final class ExperimentalFrameExtractor {
                           configuration.mediaCodecSelector,
                           videoRendererEventListener,
                           /* toneMapHdrToSdr= */ !configuration.extractHdrFrames)
-                    })
+                    },
+                mediaSourceFactory)
             .setSeekParameters(configuration.seekParameters)
             .build();
     player.addAnalyticsListener(new PlayerListener());