Skip to content

Commit

Permalink
Factor out video decoding and fix two minor issues
Browse files Browse the repository at this point in the history 
1. Not treating 0 as valid buffer index
2. Not handling the case the last frame is a comparison frame

PiperOrigin-RevId: 539607482
  • Loading branch information
@claincly @tof-tof
claincly authored and tof-tof committed Jun 12, 2023
1 parent 49b893f commit 4b1ac2f
Show file tree
Hide file tree
Showing 2 changed files with 232 additions and 195 deletions.
206 changes: 11 additions & 195 deletions libraries/transformer/src/androidTest/java/androidx/media3/transformer/SsimHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,27 +17,13 @@
package androidx.media3.transformer;

import static androidx.media3.common.util.Assertions.checkNotNull;
import static androidx.media3.common.util.Assertions.checkState;
import static androidx.media3.common.util.Assertions.checkStateNotNull;
import static androidx.media3.transformer.AndroidTestUtil.MEDIA_CODEC_PRIORITY_NON_REALTIME;
import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;

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 androidx.media3.common.MimeTypes;
import androidx.media3.common.util.ConditionVariable;
import androidx.media3.common.util.Util;
import java.io.Closeable;
import java.io.IOException;
import java.nio.ByteBuffer;

Expand All @@ -59,8 +45,8 @@ public final class SsimHelper {
/** The default comparison interval. */
public static final int DEFAULT_COMPARISON_INTERVAL = 11;

private static final int IMAGE_AVAILABLE_TIMEOUT_MS = 10_000;
private static final int DECODED_IMAGE_CHANNEL_COUNT = 3;
private static final int MAX_IMAGE_READER_IMAGES_ALLOWED = 1;

/**
* Returns the mean SSIM score between the reference and the distorted video.
Expand All @@ -77,9 +63,17 @@ public static double calculate(
Context context, String referenceVideoPath, String distortedVideoPath)
throws IOException, InterruptedException {
VideoDecodingWrapper referenceDecodingWrapper =
new VideoDecodingWrapper(context, referenceVideoPath, DEFAULT_COMPARISON_INTERVAL);
new VideoDecodingWrapper(
context,
referenceVideoPath,
DEFAULT_COMPARISON_INTERVAL,
MAX_IMAGE_READER_IMAGES_ALLOWED);
VideoDecodingWrapper distortedDecodingWrapper =
new VideoDecodingWrapper(context, distortedVideoPath, DEFAULT_COMPARISON_INTERVAL);
new VideoDecodingWrapper(
context,
distortedVideoPath,
DEFAULT_COMPARISON_INTERVAL,
MAX_IMAGE_READER_IMAGES_ALLOWED);
@Nullable byte[] referenceLumaBuffer = null;
@Nullable byte[] distortedLumaBuffer = null;
double accumulatedSsim = 0.0;
Expand Down Expand Up @@ -156,182 +150,4 @@ private static byte[] extractLumaChannelBuffer(Image image, byte[] lumaChannelBu
private SsimHelper() {
// Prevent instantiation.
}

private static final class VideoDecodingWrapper implements Closeable {
// 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 static final int MAX_IMAGES_ALLOWED = 1;

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 by
* SSIM.
* @throws IOException When failed to open the video file.
*/
public VideoDecodingWrapper(Context context, String filePath, int comparisonInterval)
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, MAX_IMAGES_ALLOWED);
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) {
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();
}
}
}
Loading

0 comments on commit 4b1ac2f

Please sign in to comment.