Factor out video decoding and fix two minor issues

1. Not treating 0 as valid buffer index
2. Not handling the case the last frame is a comparison frame

PiperOrigin-RevId: 539607482
(cherry picked from commit 750988797436a3a7af0c3ff8deb8bda22da12038)

library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/VideoDecodingWrapper.java

+/*
+ * Copyright 2023 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.android.exoplayer2.transformer;
+
+import static com.google.android.exoplayer2.transformer.AndroidTestUtil.MEDIA_CODEC_PRIORITY_NON_REALTIME;
+import static com.google.android.exoplayer2.util.Assertions.checkNotNull;
+import static com.google.android.exoplayer2.util.Assertions.checkState;
+import static com.google.android.exoplayer2.util.Assertions.checkStateNotNull;
+import static com.google.common.truth.Truth.assertThat;
+
+import android.content.Context;
+import android.content.res.AssetFileDescriptor;
+import android.graphics.ImageFormat;
+import android.media.Image;
+import android.media.ImageReader;
+import android.media.MediaCodec;
+import android.media.MediaCodecInfo;
+import android.media.MediaExtractor;
+import android.media.MediaFormat;
+import android.os.Handler;
+import androidx.annotation.Nullable;
+import com.google.android.exoplayer2.util.ConditionVariable;
+import com.google.android.exoplayer2.util.MimeTypes;
+import com.google.android.exoplayer2.util.Util;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+
+/** A wrapper for decoding a video using {@link MediaCodec}. */
+/* package */ final class VideoDecodingWrapper implements AutoCloseable {
+
+  private static final int IMAGE_AVAILABLE_TIMEOUT_MS = 10_000;
+
+  // Use ExoPlayer's 10ms timeout setting. In practise, the test durations from using timeouts of
+  // 1/10/100ms don't differ significantly.
+  private static final long DEQUEUE_TIMEOUT_US = 10_000;
+  // SSIM should be calculated using the luma (Y') channel, thus using the YUV color space.
+  private static final int IMAGE_READER_COLOR_SPACE = ImageFormat.YUV_420_888;
+  private static final int MEDIA_CODEC_COLOR_SPACE =
+      MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible;
+  private static final String ASSET_FILE_SCHEME = "asset:///";
+
+  private final MediaFormat mediaFormat;
+  private final MediaCodec mediaCodec;
+  private final MediaExtractor mediaExtractor;
+  private final MediaCodec.BufferInfo bufferInfo;
+  private final ImageReader imageReader;
+  private final ConditionVariable imageAvailableConditionVariable;
+  private final int comparisonInterval;
+
+  private boolean isCurrentFrameComparisonFrame;
+  private boolean hasReadEndOfInputStream;
+  private boolean queuedEndOfStreamToDecoder;
+  private boolean dequeuedAllDecodedFrames;
+  private boolean isCodecStarted;
+  private int dequeuedFramesCount;
+
+  /**
+   * Creates a new instance.
+   *
+   * @param context The {@link Context}.
+   * @param filePath The path to the video file.
+   * @param comparisonInterval The number of frames between the frames selected for comparison.
+   * @param maxImagesAllowed The max number of images allowed in {@link ImageReader}.
+   * @throws IOException When failed to open the video file.
+   */
+  public VideoDecodingWrapper(
+      Context context, String filePath, int comparisonInterval, int maxImagesAllowed)
+      throws IOException {
+    this.comparisonInterval = comparisonInterval;
+    mediaExtractor = new MediaExtractor();
+    bufferInfo = new MediaCodec.BufferInfo();
+
+    if (filePath.contains(ASSET_FILE_SCHEME)) {
+      AssetFileDescriptor assetFd =
+          context.getAssets().openFd(filePath.replace(ASSET_FILE_SCHEME, ""));
+      mediaExtractor.setDataSource(
+          assetFd.getFileDescriptor(), assetFd.getStartOffset(), assetFd.getLength());
+    } else {
+      mediaExtractor.setDataSource(filePath);
+    }
+
+    @Nullable MediaFormat mediaFormat = null;
+    for (int i = 0; i < mediaExtractor.getTrackCount(); i++) {
+      if (MimeTypes.isVideo(mediaExtractor.getTrackFormat(i).getString(MediaFormat.KEY_MIME))) {
+        mediaFormat = mediaExtractor.getTrackFormat(i);
+        mediaExtractor.selectTrack(i);
+        break;
+      }
+    }
+
+    checkStateNotNull(mediaFormat);
+    checkState(mediaFormat.containsKey(MediaFormat.KEY_WIDTH));
+    int width = mediaFormat.getInteger(MediaFormat.KEY_WIDTH);
+    checkState(mediaFormat.containsKey(MediaFormat.KEY_HEIGHT));
+    int height = mediaFormat.getInteger(MediaFormat.KEY_HEIGHT);
+
+    // Create a handler for the main thread to receive image available notifications. The current
+    // (test) thread blocks until this callback is received.
+    Handler mainThreadHandler = Util.createHandlerForCurrentOrMainLooper();
+    imageAvailableConditionVariable = new ConditionVariable();
+    imageReader =
+        ImageReader.newInstance(width, height, IMAGE_READER_COLOR_SPACE, maxImagesAllowed);
+    imageReader.setOnImageAvailableListener(
+        imageReader -> imageAvailableConditionVariable.open(), mainThreadHandler);
+
+    String sampleMimeType = checkNotNull(mediaFormat.getString(MediaFormat.KEY_MIME));
+    mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MEDIA_CODEC_COLOR_SPACE);
+    mediaFormat.setInteger(MediaFormat.KEY_PRIORITY, MEDIA_CODEC_PRIORITY_NON_REALTIME);
+    this.mediaFormat = mediaFormat;
+    mediaCodec = MediaCodec.createDecoderByType(sampleMimeType);
+  }
+
+  /**
+   * Returns the next decoded comparison frame, or {@code null} if the stream has ended. The caller
+   * takes ownership of any returned image and is responsible for closing it before calling this
+   * method again.
+   */
+  @Nullable
+  public Image runUntilComparisonFrameOrEnded() throws InterruptedException {
+    if (!isCodecStarted) {
+      mediaCodec.configure(
+          mediaFormat, imageReader.getSurface(), /* crypto= */ null, /* flags= */ 0);
+      mediaCodec.start();
+      isCodecStarted = true;
+    }
+    while (!hasEnded() && !isCurrentFrameComparisonFrame) {
+      while (dequeueOneFrameFromDecoder()) {}
+      while (queueOneFrameToDecoder()) {}
+    }
+    if (isCurrentFrameComparisonFrame && !hasEnded()) {
+      isCurrentFrameComparisonFrame = false;
+      assertThat(imageAvailableConditionVariable.block(IMAGE_AVAILABLE_TIMEOUT_MS)).isTrue();
+      imageAvailableConditionVariable.close();
+      return imageReader.acquireLatestImage();
+    }
+    return null;
+  }
+
+  /** Returns whether decoding has ended. */
+  private boolean hasEnded() {
+    return dequeuedAllDecodedFrames;
+  }
+
+  /** Returns whether a frame is queued to the {@link MediaCodec decoder}. */
+  private boolean queueOneFrameToDecoder() {
+    if (queuedEndOfStreamToDecoder) {
+      return false;
+    }
+
+    int inputBufferIndex = mediaCodec.dequeueInputBuffer(DEQUEUE_TIMEOUT_US);
+    if (inputBufferIndex < 0) {
+      return false;
+    }
+
+    if (hasReadEndOfInputStream) {
+      mediaCodec.queueInputBuffer(
+          inputBufferIndex,
+          /* offset= */ 0,
+          /* size= */ 0,
+          /* presentationTimeUs= */ 0,
+          MediaCodec.BUFFER_FLAG_END_OF_STREAM);
+      queuedEndOfStreamToDecoder = true;
+      return false;
+    }
+
+    ByteBuffer inputBuffer = checkNotNull(mediaCodec.getInputBuffer(inputBufferIndex));
+    int sampleSize = mediaExtractor.readSampleData(inputBuffer, /* offset= */ 0);
+    mediaCodec.queueInputBuffer(
+        inputBufferIndex,
+        /* offset= */ 0,
+        sampleSize,
+        mediaExtractor.getSampleTime(),
+        mediaExtractor.getSampleFlags());
+    // MediaExtractor.advance does not reliably return false for end-of-stream, so check sample
+    // metadata instead as a more reliable signal. See [internal: b/121204004].
+    mediaExtractor.advance();
+    hasReadEndOfInputStream = mediaExtractor.getSampleTime() == -1;
+    return true;
+  }
+
+  /** Returns whether a frame is decoded, renders the frame if the frame is a comparison frame. */
+  private boolean dequeueOneFrameFromDecoder() {
+    if (isCurrentFrameComparisonFrame) {
+      return false;
+    }
+
+    int outputBufferIndex = mediaCodec.dequeueOutputBuffer(bufferInfo, DEQUEUE_TIMEOUT_US);
+    if (outputBufferIndex < 0) {
+      return false;
+    }
+    isCurrentFrameComparisonFrame = dequeuedFramesCount % comparisonInterval == 0;
+    dequeuedFramesCount++;
+    mediaCodec.releaseOutputBuffer(outputBufferIndex, /* render= */ isCurrentFrameComparisonFrame);
+
+    if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
+      dequeuedAllDecodedFrames = true;
+    }
+    return true;
+  }
+
+  @Override
+  public void close() {
+    mediaExtractor.release();
+    mediaCodec.release();
+    imageReader.close();
+  }
+}