Commit 9e62ea3f by andrewlewis Committed by Ian Baker

Move image buffer extraction to test thread

This also ensures that if there's an error reading the image data then this
gets surfaced as an analysis exception.

PiperOrigin-RevId: 435680785
parent 12543a96
...@@ -20,6 +20,7 @@ import static androidx.media3.common.util.Assertions.checkNotNull; ...@@ -20,6 +20,7 @@ import static androidx.media3.common.util.Assertions.checkNotNull;
import static androidx.media3.common.util.Assertions.checkState; import static androidx.media3.common.util.Assertions.checkState;
import static androidx.media3.common.util.Assertions.checkStateNotNull; import static androidx.media3.common.util.Assertions.checkStateNotNull;
import static com.google.common.truth.Truth.assertThat; import static com.google.common.truth.Truth.assertThat;
import static com.google.common.truth.Truth.assertWithMessage;
import static java.lang.Math.pow; import static java.lang.Math.pow;
import android.content.Context; import android.content.Context;
...@@ -33,15 +34,12 @@ import android.media.MediaExtractor; ...@@ -33,15 +34,12 @@ import android.media.MediaExtractor;
import android.media.MediaFormat; import android.media.MediaFormat;
import android.os.Handler; import android.os.Handler;
import androidx.annotation.Nullable; import androidx.annotation.Nullable;
import androidx.media3.common.Format;
import androidx.media3.common.MimeTypes; import androidx.media3.common.MimeTypes;
import androidx.media3.common.util.ConditionVariable;
import androidx.media3.common.util.Util; import androidx.media3.common.util.Util;
import java.io.Closeable; import java.io.Closeable;
import java.io.IOException; import java.io.IOException;
import java.nio.ByteBuffer; import java.nio.ByteBuffer;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
/** /**
* A helper for calculating SSIM score for transcoded videos. * A helper for calculating SSIM score for transcoded videos.
...@@ -58,26 +56,8 @@ public final class SsimHelper { ...@@ -58,26 +56,8 @@ public final class SsimHelper {
/** The default comparison interval. */ /** The default comparison interval. */
public static final int DEFAULT_COMPARISON_INTERVAL = 11; public static final int DEFAULT_COMPARISON_INTERVAL = 11;
private static final int SURFACE_WAIT_MS = 10; private static final int IMAGE_AVAILABLE_TIMEOUT_MS = 10_000;
private static final int DECODED_IMAGE_CHANNEL_COUNT = 3; private static final int DECODED_IMAGE_CHANNEL_COUNT = 3;
private static final int[] EMPTY_BUFFER = new int[0];
private final Context context;
private final String expectedVideoPath;
private final String actualVideoPath;
private final int comparisonInterval;
private @MonotonicNonNull VideoDecodingWrapper expectedDecodingWrapper;
private @MonotonicNonNull VideoDecodingWrapper actualDecodingWrapper;
private double accumulatedSsim;
private int comparedImagesCount;
// These atomic fields are read on both test thread (where MediaCodec is controlled) and set on
// the main thread (where ImageReader invokes its callback).
private final AtomicReference<int[]> expectedLumaBuffer;
private final AtomicReference<int[]> actualLumaBuffer;
private final AtomicInteger width;
private final AtomicInteger height;
/** /**
* Returns the mean SSIM score between the expected and the actual video. * Returns the mean SSIM score between the expected and the actual video.
...@@ -92,95 +72,49 @@ public final class SsimHelper { ...@@ -92,95 +72,49 @@ public final class SsimHelper {
*/ */
public static double calculate(Context context, String expectedVideoPath, String actualVideoPath) public static double calculate(Context context, String expectedVideoPath, String actualVideoPath)
throws IOException, InterruptedException { throws IOException, InterruptedException {
return new SsimHelper(context, expectedVideoPath, actualVideoPath, DEFAULT_COMPARISON_INTERVAL) VideoDecodingWrapper expectedDecodingWrapper =
.calculateSsim(); new VideoDecodingWrapper(context, expectedVideoPath, DEFAULT_COMPARISON_INTERVAL);
} VideoDecodingWrapper actualDecodingWrapper =
new VideoDecodingWrapper(context, actualVideoPath, DEFAULT_COMPARISON_INTERVAL);
private SsimHelper( double accumulatedSsim = 0.0;
Context context, String expectedVideoPath, String actualVideoPath, int comparisonInterval) { int comparedImagesCount = 0;
this.context = context;
this.expectedVideoPath = expectedVideoPath;
this.actualVideoPath = actualVideoPath;
this.comparisonInterval = comparisonInterval;
this.expectedLumaBuffer = new AtomicReference<>(EMPTY_BUFFER);
this.actualLumaBuffer = new AtomicReference<>(EMPTY_BUFFER);
this.width = new AtomicInteger(Format.NO_VALUE);
this.height = new AtomicInteger(Format.NO_VALUE);
}
/** Calculates the SSIM score between the two videos. */
private double calculateSsim() throws InterruptedException, IOException {
// The test thread has no looper, so a handler is created on which the
// ImageReader.OnImageAvailableListener is called.
Handler mainThreadHandler = Util.createHandlerForCurrentOrMainLooper();
ImageReader.OnImageAvailableListener onImageAvailableListener = this::onImageAvailableListener;
expectedDecodingWrapper =
new VideoDecodingWrapper(
context,
expectedVideoPath,
onImageAvailableListener,
mainThreadHandler,
comparisonInterval);
actualDecodingWrapper =
new VideoDecodingWrapper(
context,
actualVideoPath,
onImageAvailableListener,
mainThreadHandler,
comparisonInterval);
try { try {
while (!expectedDecodingWrapper.hasEnded() && !actualDecodingWrapper.hasEnded()) { while (true) {
if (!expectedDecodingWrapper.runUntilComparisonFrameOrEnded() @Nullable Image expectedImage = expectedDecodingWrapper.runUntilComparisonFrameOrEnded();
|| !actualDecodingWrapper.runUntilComparisonFrameOrEnded()) { @Nullable Image actualImage = actualDecodingWrapper.runUntilComparisonFrameOrEnded();
continue; if (expectedImage == null) {
assertThat(actualImage).isNull();
break;
} }
checkNotNull(actualImage);
while (expectedLumaBuffer.get() == EMPTY_BUFFER || actualLumaBuffer.get() == EMPTY_BUFFER) { int width = expectedImage.getWidth();
// Wait for the ImageReader to call onImageAvailable and process the luma channel on the int height = expectedImage.getHeight();
// main thread. assertThat(actualImage.getWidth()).isEqualTo(width);
Thread.sleep(SURFACE_WAIT_MS); assertThat(actualImage.getHeight()).isEqualTo(height);
} try {
accumulatedSsim += accumulatedSsim +=
SsimCalculator.calculate( SsimCalculator.calculate(
expectedLumaBuffer.get(), extractLumaChannelBuffer(expectedImage),
actualLumaBuffer.get(), extractLumaChannelBuffer(actualImage),
/* offset= */ 0, /* offset= */ 0,
/* stride= */ width.get(), /* stride= */ width,
width.get(), width,
height.get()); height);
} finally {
expectedImage.close();
actualImage.close();
}
comparedImagesCount++; comparedImagesCount++;
expectedLumaBuffer.set(EMPTY_BUFFER);
actualLumaBuffer.set(EMPTY_BUFFER);
} }
} finally { } finally {
expectedDecodingWrapper.close(); expectedDecodingWrapper.close();
actualDecodingWrapper.close(); actualDecodingWrapper.close();
} }
assertWithMessage("Input had no frames.").that(comparedImagesCount).isGreaterThan(0);
if (comparedImagesCount == 0) {
throw new IOException("Input had no frames.");
}
return accumulatedSsim / comparedImagesCount; return accumulatedSsim / comparedImagesCount;
} }
private void onImageAvailableListener(ImageReader imageReader) {
// This method is invoked on the main thread.
Image image = imageReader.acquireLatestImage();
int[] lumaBuffer = extractLumaChannelBuffer(image);
width.set(image.getWidth());
height.set(image.getHeight());
image.close();
if (imageReader == checkNotNull(expectedDecodingWrapper).imageReader) {
expectedLumaBuffer.set(lumaBuffer);
} else if (imageReader == checkNotNull(actualDecodingWrapper).imageReader) {
actualLumaBuffer.set(lumaBuffer);
} else {
throw new IllegalStateException("Unexpected ImageReader.");
}
}
/** /**
* Returns the buffer of the luma (Y) channel of the image. * Returns the buffer of the luma (Y) channel of the image.
* *
...@@ -206,6 +140,10 @@ public final class SsimHelper { ...@@ -206,6 +140,10 @@ public final class SsimHelper {
return lumaChannelBuffer; return lumaChannelBuffer;
} }
private SsimHelper() {
// Prevent instantiation.
}
private static final class VideoDecodingWrapper implements Closeable { private static final class VideoDecodingWrapper implements Closeable {
// Use ExoPlayer's 10ms timeout setting. In practise, the test durations from using timeouts of // Use ExoPlayer's 10ms timeout setting. In practise, the test durations from using timeouts of
// 1/10/100ms don't differ significantly. // 1/10/100ms don't differ significantly.
...@@ -221,6 +159,7 @@ public final class SsimHelper { ...@@ -221,6 +159,7 @@ public final class SsimHelper {
private final MediaExtractor mediaExtractor; private final MediaExtractor mediaExtractor;
private final MediaCodec.BufferInfo bufferInfo; private final MediaCodec.BufferInfo bufferInfo;
private final ImageReader imageReader; private final ImageReader imageReader;
private final ConditionVariable imageAvailableConditionVariable;
private final int comparisonInterval; private final int comparisonInterval;
private boolean isCurrentFrameComparisonFrame; private boolean isCurrentFrameComparisonFrame;
...@@ -234,18 +173,11 @@ public final class SsimHelper { ...@@ -234,18 +173,11 @@ public final class SsimHelper {
* *
* @param context The {@link Context}. * @param context The {@link Context}.
* @param filePath The path to the video file. * @param filePath The path to the video file.
* @param imageAvailableListener An {@link ImageReader.OnImageAvailableListener} implementation.
* @param handler The {@link Handler} on which the {@code imageAvailableListener} is called.
* @param comparisonInterval The number of frames between the frames selected for comparison by * @param comparisonInterval The number of frames between the frames selected for comparison by
* SSIM. * SSIM.
* @throws IOException When failed to open the video file. * @throws IOException When failed to open the video file.
*/ */
public VideoDecodingWrapper( public VideoDecodingWrapper(Context context, String filePath, int comparisonInterval)
Context context,
String filePath,
ImageReader.OnImageAvailableListener imageAvailableListener,
Handler handler,
int comparisonInterval)
throws IOException { throws IOException {
this.comparisonInterval = comparisonInterval; this.comparisonInterval = comparisonInterval;
mediaExtractor = new MediaExtractor(); mediaExtractor = new MediaExtractor();
...@@ -275,9 +207,14 @@ public final class SsimHelper { ...@@ -275,9 +207,14 @@ public final class SsimHelper {
checkState(mediaFormat.containsKey(MediaFormat.KEY_HEIGHT)); checkState(mediaFormat.containsKey(MediaFormat.KEY_HEIGHT));
int height = mediaFormat.getInteger(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 =
ImageReader.newInstance(width, height, IMAGE_READER_COLOR_SPACE, MAX_IMAGES_ALLOWED); ImageReader.newInstance(width, height, IMAGE_READER_COLOR_SPACE, MAX_IMAGES_ALLOWED);
imageReader.setOnImageAvailableListener(imageAvailableListener, handler); imageReader.setOnImageAvailableListener(
imageReader -> imageAvailableConditionVariable.open(), mainThreadHandler);
String sampleMimeType = checkNotNull(mediaFormat.getString(MediaFormat.KEY_MIME)); String sampleMimeType = checkNotNull(mediaFormat.getString(MediaFormat.KEY_MIME));
mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MEDIA_CODEC_COLOR_SPACE); mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MEDIA_CODEC_COLOR_SPACE);
...@@ -288,29 +225,28 @@ public final class SsimHelper { ...@@ -288,29 +225,28 @@ public final class SsimHelper {
} }
/** /**
* Run decoding until a comparison frame is rendered, or decoding has ended. * 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
* <p>The method returns after rendering the comparison frame. There is no guarantee that the * this method again.
* frame is available for processing at this time.
*
* @return {@code true} when a comparison frame is encountered, or {@code false} if decoding
* {@link #hasEnded() had ended}.
*/ */
public boolean runUntilComparisonFrameOrEnded() { @Nullable
public Image runUntilComparisonFrameOrEnded() throws InterruptedException {
while (!hasEnded() && !isCurrentFrameComparisonFrame) { while (!hasEnded() && !isCurrentFrameComparisonFrame) {
while (dequeueOneFrameFromDecoder()) {} while (dequeueOneFrameFromDecoder()) {}
while (queueOneFrameToDecoder()) {} while (queueOneFrameToDecoder()) {}
} }
if (isCurrentFrameComparisonFrame) { if (isCurrentFrameComparisonFrame) {
isCurrentFrameComparisonFrame = false; isCurrentFrameComparisonFrame = false;
return true; assertThat(imageAvailableConditionVariable.block(IMAGE_AVAILABLE_TIMEOUT_MS)).isTrue();
imageAvailableConditionVariable.close();
return imageReader.acquireLatestImage();
} }
return false; return null;
} }
/** Returns whether decoding has ended. */ /** Returns whether decoding has ended. */
public boolean hasEnded() { private boolean hasEnded() {
return queuedEndOfStreamToDecoder && dequeuedAllDecodedFrames; return dequeuedAllDecodedFrames;
} }
/** Returns whether a frame is queued to the {@link MediaCodec decoder}. */ /** Returns whether a frame is queued to the {@link MediaCodec decoder}. */
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment