Add ability to queue bitmap by timestamp to the sampleConsumer

PiperOrigin-RevId: 557837922

libraries/transformer/src/main/java/androidx/media3/transformer/SampleConsumer.java

@@ -22,6 +22,7 @@ import androidx.media3.common.ColorInfo;
 import androidx.media3.common.OnInputFrameProcessedListener;
 import androidx.media3.common.util.UnstableApi;
 import androidx.media3.decoder.DecoderInputBuffer;
+import java.util.Iterator;
 
 /** Consumer of encoded media samples, raw audio or raw video frames. */
 @UnstableApi
@@ -77,10 +78,25 @@ public interface SampleConsumer {
    * @return Whether the {@link Bitmap} was successfully queued. If {@code false}, the caller should
    *     try again later.
    */
+  // TODO(b/262693274): Delete this method and usages in favor of the one below (Note it is not
+  //   deprecated because transformer still relies on this method for frame duplication).
   default boolean queueInputBitmap(Bitmap inputBitmap, long durationUs, int frameRate) {
     throw new UnsupportedOperationException();
   }
 
+  /**
+   * Attempts to provide an input {@link Bitmap} to the consumer.
+   *
+   * <p>Should only be used for image data.
+   *
+   * @param inputBitmap The {@link Bitmap} to queue to the consumer.
+   * @param inStreamOffsetsUs The times within the current stream that the bitmap should be
+   *     displayed at. The timestamps should be monotonically increasing.
+   */
+  default boolean queueInputBitmap(Bitmap inputBitmap, Iterator<Long> inStreamOffsetsUs) {
+    throw new UnsupportedOperationException();
+  }
+
   // Methods to pass raw video input.
 
   /**

libraries/transformer/src/main/java/androidx/media3/transformer/SequenceAssetLoader.java

@@ -39,6 +39,7 @@ import androidx.media3.decoder.DecoderInputBuffer;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import java.util.HashMap;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -347,6 +348,38 @@ import java.util.concurrent.atomic.AtomicInteger;
     sequenceAssetLoaderListener.onError(exportException);
   }
 
+  /**
+   * Given an {@link Iterator}, creates an iterator that includes all the values in the original
+   * iterator (in the same order) up to and including the first occurrence of the {@code
+   * clippingValue}.
+   */
+  private static final class ClippingIterator implements Iterator<Long> {
+
+    private final Iterator<Long> iterator;
+    private final long clippingValue;
+    private boolean hasReachedClippingValue;
+
+    public ClippingIterator(Iterator<Long> iterator, long clippingValue) {
+      this.iterator = iterator;
+      this.clippingValue = clippingValue;
+    }
+
+    @Override
+    public boolean hasNext() {
+      return !hasReachedClippingValue && iterator.hasNext();
+    }
+
+    @Override
+    public Long next() {
+      checkState(hasNext());
+      Long next = iterator.next();
+      if (clippingValue == next) {
+        hasReachedClippingValue = true;
+      }
+      return next;
+    }
+  }
+
   // Classes accessed from AssetLoader threads.
 
   private final class SampleConsumerWrapper implements SampleConsumer {
@@ -396,8 +429,6 @@ import java.util.concurrent.atomic.AtomicInteger;
       return true;
     }
 
-    // TODO(b/262693274): Test that concatenate 2 images or an image and a video works as expected
-    //  once ImageAssetLoader implementation is complete.
     @Override
     public boolean queueInputBitmap(Bitmap inputBitmap, long durationUs, int frameRate) {
       if (isLooping && totalDurationUs + durationUs > maxSequenceDurationUs) {
@@ -418,6 +449,33 @@ import java.util.concurrent.atomic.AtomicInteger;
       return sampleConsumer.queueInputBitmap(inputBitmap, durationUs, frameRate);
     }
 
+    @Override
+    public boolean queueInputBitmap(Bitmap inputBitmap, Iterator<Long> inStreamOffsetsUs) {
+      Iterator<Long> iteratorToUse = inStreamOffsetsUs;
+      if (isLooping) {
+        long durationLeftUs = maxSequenceDurationUs - totalDurationUs;
+        if (durationLeftUs <= 0) {
+          if (!videoLoopingEnded) {
+            videoLoopingEnded = true;
+            signalEndOfVideoInput();
+          }
+          return false;
+        }
+        while (inStreamOffsetsUs.hasNext()) {
+          long offsetUs = inStreamOffsetsUs.next();
+          if (totalDurationUs + offsetUs > maxSequenceDurationUs) {
+            if (!isMaxSequenceDurationUsFinal) {
+              return false;
+            }
+            iteratorToUse = new ClippingIterator(inStreamOffsetsUs, offsetUs);
+            videoLoopingEnded = true;
+            break;
+          }
+        }
+      }
+      return sampleConsumer.queueInputBitmap(inputBitmap, iteratorToUse);
+    }
+
     @Override
     public void setOnInputFrameProcessedListener(OnInputFrameProcessedListener listener) {
       sampleConsumer.setOnInputFrameProcessedListener(listener);

libraries/transformer/src/main/java/androidx/media3/transformer/SingleInputVideoGraph.java

@@ -39,6 +39,7 @@ import androidx.media3.common.util.Consumer;
 import androidx.media3.common.util.Size;
 import androidx.media3.effect.Presentation;
 import com.google.common.collect.ImmutableList;
+import java.util.Iterator;
 import java.util.List;
 import java.util.concurrent.Executor;
 import java.util.concurrent.atomic.AtomicLong;
@@ -173,6 +174,12 @@ import java.util.concurrent.atomic.AtomicLong;
     return true;
   }
 
+  @Override
+  public boolean queueInputBitmap(Bitmap inputBitmap, Iterator<Long> inStreamOffsetsUs) {
+    videoFrameProcessor.queueInputBitmap(inputBitmap, inStreamOffsetsUs);
+    return true;
+  }
+
   @Override
   public void setOnInputFrameProcessedListener(OnInputFrameProcessedListener listener) {
     videoFrameProcessor.setOnInputFrameProcessedListener(listener);