From 23665090ed5cc48ac03fd9099ef1eefa2048a6d8 Mon Sep 17 00:00:00 2001
From: tofunmi <tofunmi@google.com>
Date: Fri, 8 Sep 2023 12:12:15 -0700
Subject: [PATCH] Support Jpeg image track extraction in exoplayer

PiperOrigin-RevId: 563818198
---
 .../source/DefaultMediaSourceFactory.java       | 10 ++++++++++
 .../exoplayer/e2etest/ImagePlaybackTest.java    |  5 +++--
 .../extractor/DefaultExtractorsFactory.java     | 17 ++++++++++++++++-
 .../jpeg/non-motion-photo-shortened.jpg.dump    |  5 +++++
 4 files changed, 34 insertions(+), 3 deletions(-)
 create mode 100644 libraries/test_data/src/test/assets/playbackdumps/jpeg/non-motion-photo-shortened.jpg.dump

diff --git a/libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/DefaultMediaSourceFactory.java b/libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/DefaultMediaSourceFactory.java
index eb7a8f0524..0c0b9c3a09 100644
--- a/libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/DefaultMediaSourceFactory.java
+++ b/libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/DefaultMediaSourceFactory.java
@@ -48,6 +48,7 @@ import androidx.media3.extractor.ExtractorsFactory;
 import androidx.media3.extractor.PositionHolder;
 import androidx.media3.extractor.SeekMap;
 import androidx.media3.extractor.TrackOutput;
+import androidx.media3.extractor.jpeg.JpegExtractor;
 import androidx.media3.extractor.text.DefaultSubtitleParserFactory;
 import androidx.media3.extractor.text.SubtitleExtractor;
 import com.google.common.base.Supplier;
@@ -438,6 +439,9 @@ public final class DefaultMediaSourceFactory implements MediaSourceFactory {
     int type =
         Util.inferContentTypeForUriAndMimeType(
             mediaItem.localConfiguration.uri, mediaItem.localConfiguration.mimeType);
+    if (mediaItem.localConfiguration.imageDurationMs != C.TIME_UNSET) {
+      delegateFactoryLoader.setJpegExtractorFlags(JpegExtractor.FLAG_READ_IMAGE);
+    }
     @Nullable
     MediaSource.Factory mediaSourceFactory = delegateFactoryLoader.getMediaSourceFactory(type);
     checkStateNotNull(
@@ -650,6 +654,12 @@ public final class DefaultMediaSourceFactory implements MediaSourceFactory {
       }
     }
 
+    public void setJpegExtractorFlags(@JpegExtractor.Flags int flags) {
+      if (this.extractorsFactory instanceof DefaultExtractorsFactory) {
+        ((DefaultExtractorsFactory) this.extractorsFactory).setJpegExtractorFlags(flags);
+      }
+    }
+
     private void ensureAllSuppliersAreLoaded() {
       maybeLoadSupplier(C.CONTENT_TYPE_DASH);
       maybeLoadSupplier(C.CONTENT_TYPE_SS);
diff --git a/libraries/exoplayer/src/test/java/androidx/media3/exoplayer/e2etest/ImagePlaybackTest.java b/libraries/exoplayer/src/test/java/androidx/media3/exoplayer/e2etest/ImagePlaybackTest.java
index 7b1de64809..eb03d2c571 100644
--- a/libraries/exoplayer/src/test/java/androidx/media3/exoplayer/e2etest/ImagePlaybackTest.java
+++ b/libraries/exoplayer/src/test/java/androidx/media3/exoplayer/e2etest/ImagePlaybackTest.java
@@ -50,8 +50,9 @@ public class ImagePlaybackTest {
     // TODO(b/289989736): When extraction for other types of images is implemented, add those image
     //   types to this list.
     // Robolectric's NativeShadowBitmapFactory doesn't support decoding HEIF format, so we don't
-    // test that here.
-    return ImmutableList.of("png/non-motion-photo-shortened.png");
+    // test that format here.
+    return ImmutableList.of(
+        "png/non-motion-photo-shortened.png", "jpeg/non-motion-photo-shortened.jpg");
   }
 
   @Test
diff --git a/libraries/extractor/src/main/java/androidx/media3/extractor/DefaultExtractorsFactory.java b/libraries/extractor/src/main/java/androidx/media3/extractor/DefaultExtractorsFactory.java
index 86c6fc5a27..d1f6b46bce 100644
--- a/libraries/extractor/src/main/java/androidx/media3/extractor/DefaultExtractorsFactory.java
+++ b/libraries/extractor/src/main/java/androidx/media3/extractor/DefaultExtractorsFactory.java
@@ -153,6 +153,7 @@ public final class DefaultExtractorsFactory implements ExtractorsFactory {
   private int tsTimestampSearchBytes;
   private boolean textTrackTranscodingEnabled;
   private SubtitleParser.Factory subtitleParserFactory;
+  private @JpegExtractor.Flags int jpegFlags;
 
   public DefaultExtractorsFactory() {
     tsMode = TsExtractor.MODE_SINGLE_PMT;
@@ -391,6 +392,20 @@ public final class DefaultExtractorsFactory implements ExtractorsFactory {
     return this;
   }
 
+  /**
+   * Sets flags for {@link JpegExtractor} instances created by the factory.
+   *
+   * @see JpegExtractor#JpegExtractor(int)
+   * @param flags The flags to use.
+   * @return The factory, for convenience.
+   */
+  @CanIgnoreReturnValue
+  public synchronized DefaultExtractorsFactory setJpegExtractorFlags(
+      @JpegExtractor.Flags int flags) {
+    this.jpegFlags = flags;
+    return this;
+  }
+
   @Override
   public synchronized Extractor[] createExtractors() {
     return createExtractors(Uri.EMPTY, new HashMap<>());
@@ -509,7 +524,7 @@ public final class DefaultExtractorsFactory implements ExtractorsFactory {
         extractors.add(new WavExtractor());
         break;
       case FileTypes.JPEG:
-        extractors.add(new JpegExtractor());
+        extractors.add(new JpegExtractor(jpegFlags));
         break;
       case FileTypes.MIDI:
         @Nullable Extractor midiExtractor = MIDI_EXTENSION_LOADER.getExtractor();
diff --git a/libraries/test_data/src/test/assets/playbackdumps/jpeg/non-motion-photo-shortened.jpg.dump b/libraries/test_data/src/test/assets/playbackdumps/jpeg/non-motion-photo-shortened.jpg.dump
new file mode 100644
index 0000000000..b6ff44ed89
--- /dev/null
+++ b/libraries/test_data/src/test/assets/playbackdumps/jpeg/non-motion-photo-shortened.jpg.dump
@@ -0,0 +1,5 @@
+ImageOutput:
+  rendered image count = 1
+  image output #1:
+    presentationTimeUs = 0
+    bitmap hash = 1655078913