Update VideoSamplePipeline to take account of decoder colors

PiperOrigin-RevId: 526940261

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

@@ -16,9 +16,11 @@
 
 package androidx.media3.transformer;
 
+import static androidx.media3.common.ColorInfo.SDR_BT709_LIMITED;
 import static androidx.media3.common.ColorInfo.isTransferHdr;
 import static androidx.media3.common.VideoFrameProcessor.INPUT_TYPE_BITMAP;
 import static androidx.media3.common.VideoFrameProcessor.INPUT_TYPE_SURFACE;
+import static androidx.media3.common.util.Assertions.checkArgument;
 import static androidx.media3.common.util.Assertions.checkNotNull;
 import static androidx.media3.transformer.EncoderUtil.getSupportedEncodersForHdrEditing;
 import static androidx.media3.transformer.TransformationRequest.HDR_MODE_KEEP_HDR;
@@ -59,9 +61,7 @@ import org.checkerframework.dataflow.qual.Pure;
 /** Pipeline to process, re-encode and mux raw video frames. */
 /* package */ final class VideoSamplePipeline extends SamplePipeline {
 
-  /** MIME type to use for output video if the input type is not a video. */
+  private static final String TAG = "VideoSamplePipeline";
-  private static final String DEFAULT_OUTPUT_MIME_TYPE = MimeTypes.VIDEO_H265;
-
   private final AtomicLong mediaItemOffsetUs;
   private final VideoFrameProcessor videoFrameProcessor;
   private final ColorInfo videoFrameProcessorInputColor;
@@ -99,22 +99,30 @@ import org.checkerframework.dataflow.qual.Pure;
     encoderOutputBuffer =
         new DecoderInputBuffer(DecoderInputBuffer.BUFFER_REPLACEMENT_MODE_DISABLED);
 
+    ColorInfo decoderInputColor;
+    if (firstInputFormat.colorInfo == null || !firstInputFormat.colorInfo.isValid()) {
+      Log.d(TAG, "colorInfo is null or invalid. Defaulting to SDR_BT709_LIMITED.");
+      decoderInputColor = ColorInfo.SDR_BT709_LIMITED;
+    } else {
+      decoderInputColor = firstInputFormat.colorInfo;
+    }
+
     encoderWrapper =
         new EncoderWrapper(
             encoderFactory,
-            firstInputFormat,
+            firstInputFormat.buildUpon().setColorInfo(decoderInputColor).build(),
             muxerWrapper.getSupportedSampleMimeTypes(C.TRACK_TYPE_VIDEO),
             transformationRequest,
             fallbackListener);
 
-    ColorInfo encoderInputColor = encoderWrapper.getSupportedInputColor();
+    boolean isMediaCodecToneMapping =
-    // If not tone mapping using OpenGL, the decoder will output the encoderInputColor,
+        encoderWrapper.getSupportedInputHdrMode() == HDR_MODE_TONE_MAP_HDR_TO_SDR_USING_MEDIACODEC
-    // possibly by tone mapping.
+            && ColorInfo.isTransferHdr(decoderInputColor);
+    videoFrameProcessorInputColor = isMediaCodecToneMapping ? SDR_BT709_LIMITED : decoderInputColor;
+
     boolean isGlToneMapping =
-        ColorInfo.isTransferHdr(firstInputFormat.colorInfo)
+        ColorInfo.isTransferHdr(decoderInputColor)
             && transformationRequest.hdrMode == HDR_MODE_TONE_MAP_HDR_TO_SDR_USING_OPEN_GL;
-    videoFrameProcessorInputColor =
-        isGlToneMapping ? checkNotNull(firstInputFormat.colorInfo) : encoderInputColor;
     // For consistency with the Android platform, OpenGL tone mapping outputs colors with
     // C.COLOR_TRANSFER_GAMMA_2_2 instead of C.COLOR_TRANSFER_SDR, and outputs this as
     // C.COLOR_TRANSFER_SDR to the encoder.
@@ -125,7 +133,7 @@ import org.checkerframework.dataflow.qual.Pure;
                 .setColorRange(C.COLOR_RANGE_LIMITED)
                 .setColorTransfer(C.COLOR_TRANSFER_GAMMA_2_2)
                 .build()
-            : encoderInputColor;
+            : videoFrameProcessorInputColor;
     List<Effect> effectsWithPresentation = new ArrayList<>(effects);
     if (presentation != null) {
       effectsWithPresentation.add(presentation);
@@ -339,7 +347,8 @@ import org.checkerframework.dataflow.qual.Pure;
    */
   @VisibleForTesting
   /* package */ static final class EncoderWrapper {
-    private static final String TAG = "EncoderWrapper";
+    /** MIME type to use for output video if the input type is not a video. */
+    private static final String DEFAULT_OUTPUT_MIME_TYPE = MimeTypes.VIDEO_H265;
 
     private final Codec.EncoderFactory encoderFactory;
     private final Format inputFormat;
@@ -347,7 +356,8 @@ import org.checkerframework.dataflow.qual.Pure;
     private final TransformationRequest transformationRequest;
     private final FallbackListener fallbackListener;
     private final String requestedOutputMimeType;
-    private final boolean isHdrEditingEnabled;
+    private final boolean isInputToneMapped;
+    private final @TransformationRequest.HdrMode int supportedFallbackHdrMode;
 
     private @MonotonicNonNull SurfaceInfo encoderSurfaceInfo;
 
@@ -361,6 +371,7 @@ import org.checkerframework.dataflow.qual.Pure;
         List<String> muxerSupportedMimeTypes,
         TransformationRequest transformationRequest,
         FallbackListener fallbackListener) {
+      checkArgument(inputFormat.colorInfo != null);
       this.encoderFactory = encoderFactory;
       this.inputFormat = inputFormat;
       this.muxerSupportedMimeTypes = muxerSupportedMimeTypes;
@@ -376,26 +387,33 @@ import org.checkerframework.dataflow.qual.Pure;
         requestedOutputMimeType = inputSampleMimeType;
       }
 
-      isHdrEditingEnabled =
+      isInputToneMapped =
-          transformationRequest.hdrMode == HDR_MODE_KEEP_HDR
+          isTransferHdr(inputFormat.colorInfo)
-              && !getSupportedEncodersForHdrEditing(requestedOutputMimeType, inputFormat.colorInfo)
+              && getSupportedEncodersForHdrEditing(requestedOutputMimeType, inputFormat.colorInfo)
                   .isEmpty();
+
+      // HdrMode fallback is only supported from HDR_MODE_KEEP_HDR to
+      // HDR_MODE_TONE_MAP_HDR_TO_SDR_USING_MEDIACODEC.
+      boolean fallbackToMediaCodec =
+          isInputToneMapped && transformationRequest.hdrMode == HDR_MODE_KEEP_HDR;
+      supportedFallbackHdrMode =
+          fallbackToMediaCodec
+              ? HDR_MODE_TONE_MAP_HDR_TO_SDR_USING_MEDIACODEC
+              : transformationRequest.hdrMode;
     }
 
     /** Returns the {@link ColorInfo} expected from the input surface. */
     public ColorInfo getSupportedInputColor() {
-
-      boolean isInputToneMapped = !isHdrEditingEnabled && isTransferHdr(inputFormat.colorInfo);
       if (isInputToneMapped) {
         // When tone-mapping HDR to SDR is enabled, assume we get BT.709 to avoid having the encoder
         // populate default color info, which depends on the resolution.
         return ColorInfo.SDR_BT709_LIMITED;
       }
-      if (inputFormat.colorInfo == null || !inputFormat.colorInfo.isValid()) {
+      return checkNotNull(inputFormat.colorInfo);
-        Log.d(TAG, "colorInfo is null or invalid. Defaulting to SDR_BT709_LIMITED.");
-        return ColorInfo.SDR_BT709_LIMITED;
     }
-      return inputFormat.colorInfo;
+
+    public @TransformationRequest.HdrMode int getSupportedInputHdrMode() {
+      return supportedFallbackHdrMode;
     }
 
     @Nullable
@@ -440,17 +458,6 @@ import org.checkerframework.dataflow.qual.Pure;
 
       Format actualEncoderFormat = encoder.getConfigurationFormat();
 
-      boolean isInputToneMapped =
-          isTransferHdr(inputFormat.colorInfo) && !isTransferHdr(requestedEncoderFormat.colorInfo);
-
-      // HdrMode fallback is only supported from HDR_MODE_KEEP_HDR to
-      // HDR_MODE_TONE_MAP_HDR_TO_SDR_USING_MEDIACODEC.
-      @TransformationRequest.HdrMode
-      int supportedFallbackHdrMode =
-          isInputToneMapped && transformationRequest.hdrMode == HDR_MODE_KEEP_HDR
-              ? HDR_MODE_TONE_MAP_HDR_TO_SDR_USING_MEDIACODEC
-              : transformationRequest.hdrMode;
-
       fallbackListener.onTransformationRequestFinalized(
           createSupportedTransformationRequest(
               transformationRequest,

libraries/transformer/src/test/java/androidx/media3/transformer/VideoEncoderWrapperTest.java

@@ -24,6 +24,7 @@ import android.media.MediaFormat;
 import android.net.Uri;
 import android.os.Looper;
 import androidx.media3.common.C;
+import androidx.media3.common.ColorInfo;
 import androidx.media3.common.Format;
 import androidx.media3.common.MediaItem;
 import androidx.media3.common.MimeTypes;
@@ -63,7 +64,10 @@ public final class VideoEncoderWrapperTest {
   private final VideoSamplePipeline.EncoderWrapper encoderWrapper =
       new VideoSamplePipeline.EncoderWrapper(
           fakeEncoderFactory,
-          /* inputFormat= */ new Format.Builder().setSampleMimeType(MimeTypes.VIDEO_H264).build(),
+          /* inputFormat= */ new Format.Builder()
+              .setSampleMimeType(MimeTypes.VIDEO_H264)
+              .setColorInfo(ColorInfo.SDR_BT709_LIMITED)
+              .build(),
           /* muxerSupportedMimeTypes= */ ImmutableList.of(MimeTypes.VIDEO_H264),
           emptyTransformationRequest,
           fallbackListener);