From 7c8f6d5d69652f756c8aa4301d442b9be3890192 Mon Sep 17 00:00:00 2001
From: andrewlewis <andrewlewis@google.com>
Date: Tue, 1 Feb 2022 10:24:43 +0000
Subject: [PATCH] Add support for experimenting with HDR

- Add a checkbox in the demo app to enable experimental HDR editing.
- Add an `experimental_` method to `TransformationRequest` to enable HDR editing.
- Add fragment/vertex shaders for the experimental HDR pipeline. The main difference compared to the existing shaders is that we sample from the decoder in YUV rather than RGB (because the YUV -> RGB conversion in the graphics driver is not precisely defined, so we need to do this to get consistent results), which requires the use of ES 3, and then do a crude YUV -> RGB conversion in the shader (ignoring the input color primaries for now).
- When HDR editing is enabled, we force using `FrameEditor` (no passthrough) to avoid the need to select another edit operation, and use the new shaders. The `EGLContext` and `EGLSurface` also need to be set up differently for this path.

PiperOrigin-RevId: 425570639
---
 .../ConfigurationActivity.java                |  21 +++-
 .../transformerdemo/TransformerActivity.java  |   2 +
 .../res/layout/configuration_activity.xml     |  10 ++
 .../src/main/res/values/strings.xml           |   1 +
 .../android/exoplayer2/util/GlUtil.java       | 112 ++++++++++++++----
 .../FrameEditorDataProcessingTest.java        |   1 +
 .../transformer/FrameEditorTest.java          |   4 +-
 ...lsl => fragment_shader_copy_external.glsl} |   0
 ...fragment_shader_copy_external_yuv_es3.glsl |  29 +++++
 ...glsl => vertex_shader_transformation.glsl} |   0
 .../vertex_shader_transformation_es3.glsl     |  23 ++++
 .../exoplayer2/transformer/FrameEditor.java   |  71 +++++++++--
 .../transformer/TransformationRequest.java    |  40 ++++++-
 .../transformer/TransformerVideoRenderer.java |   3 +
 .../VideoTranscodingSamplePipeline.java       |   4 +-
 15 files changed, 274 insertions(+), 47 deletions(-)
 rename library/transformer/src/main/assets/shaders/{fragment_shader.glsl => fragment_shader_copy_external.glsl} (100%)
 create mode 100644 library/transformer/src/main/assets/shaders/fragment_shader_copy_external_yuv_es3.glsl
 rename library/transformer/src/main/assets/shaders/{vertex_shader.glsl => vertex_shader_transformation.glsl} (100%)
 create mode 100644 library/transformer/src/main/assets/shaders/vertex_shader_transformation_es3.glsl

diff --git a/demos/transformer/src/main/java/com/google/android/exoplayer2/transformerdemo/ConfigurationActivity.java b/demos/transformer/src/main/java/com/google/android/exoplayer2/transformerdemo/ConfigurationActivity.java
index a0f96d8b69..3381935df4 100644
--- a/demos/transformer/src/main/java/com/google/android/exoplayer2/transformerdemo/ConfigurationActivity.java
+++ b/demos/transformer/src/main/java/com/google/android/exoplayer2/transformerdemo/ConfigurationActivity.java
@@ -54,6 +54,7 @@ public final class ConfigurationActivity extends AppCompatActivity {
   public static final String SCALE_X = "scale_x";
   public static final String SCALE_Y = "scale_y";
   public static final String ROTATE_DEGREES = "rotate_degrees";
+  public static final String ENABLE_HDR_EDITING = "enable_hdr_editing";
   private static final String[] INPUT_URIS = {
     "https://html5demos.com/assets/dizzy.mp4",
     "https://storage.googleapis.com/exoplayer-test-media-0/android-block-1080-hevc.mp4",
@@ -69,6 +70,7 @@ public final class ConfigurationActivity extends AppCompatActivity {
   private static final String SAME_AS_INPUT_OPTION = "same as input";
 
   private @MonotonicNonNull Button chooseFileButton;
+  private @MonotonicNonNull TextView chosenFileTextView;
   private @MonotonicNonNull CheckBox removeAudioCheckbox;
   private @MonotonicNonNull CheckBox removeVideoCheckbox;
   private @MonotonicNonNull CheckBox flattenForSlowMotionCheckbox;
@@ -78,7 +80,7 @@ public final class ConfigurationActivity extends AppCompatActivity {
   private @MonotonicNonNull Spinner translateSpinner;
   private @MonotonicNonNull Spinner scaleSpinner;
   private @MonotonicNonNull Spinner rotateSpinner;
-  private @MonotonicNonNull TextView chosenFileTextView;
+  private @MonotonicNonNull CheckBox enableHdrEditingCheckBox;
   private int inputUriPosition;
 
   @Override
@@ -151,6 +153,8 @@ public final class ConfigurationActivity extends AppCompatActivity {
     rotateSpinner = findViewById(R.id.rotate_spinner);
     rotateSpinner.setAdapter(rotateAdapter);
     rotateAdapter.addAll(SAME_AS_INPUT_OPTION, "0", "10", "45", "90", "180");
+
+    enableHdrEditingCheckBox = findViewById(R.id.hdr_editing_checkbox);
   }
 
   @Override
@@ -178,7 +182,8 @@ public final class ConfigurationActivity extends AppCompatActivity {
     "resolutionHeightSpinner",
     "translateSpinner",
     "scaleSpinner",
-    "rotateSpinner"
+    "rotateSpinner",
+    "enableHdrEditingCheckBox"
   })
   private void startTransformation(View view) {
     Intent transformerIntent = new Intent(this, TransformerActivity.class);
@@ -216,6 +221,7 @@ public final class ConfigurationActivity extends AppCompatActivity {
     if (!SAME_AS_INPUT_OPTION.equals(selectedRotate)) {
       bundle.putFloat(ROTATE_DEGREES, Float.parseFloat(selectedRotate));
     }
+    bundle.putBoolean(ENABLE_HDR_EDITING, enableHdrEditingCheckBox.isChecked());
     transformerIntent.putExtras(bundle);
 
     @Nullable Uri intentUri = getIntent().getData();
@@ -247,7 +253,8 @@ public final class ConfigurationActivity extends AppCompatActivity {
     "resolutionHeightSpinner",
     "translateSpinner",
     "scaleSpinner",
-    "rotateSpinner"
+    "rotateSpinner",
+    "enableHdrEditingCheckBox"
   })
   private void onRemoveAudio(View view) {
     if (((CheckBox) view).isChecked()) {
@@ -265,7 +272,8 @@ public final class ConfigurationActivity extends AppCompatActivity {
     "resolutionHeightSpinner",
     "translateSpinner",
     "scaleSpinner",
-    "rotateSpinner"
+    "rotateSpinner",
+    "enableHdrEditingCheckBox"
   })
   private void onRemoveVideo(View view) {
     if (((CheckBox) view).isChecked()) {
@@ -282,7 +290,8 @@ public final class ConfigurationActivity extends AppCompatActivity {
     "resolutionHeightSpinner",
     "translateSpinner",
     "scaleSpinner",
-    "rotateSpinner"
+    "rotateSpinner",
+    "enableHdrEditingCheckBox"
   })
   private void enableTrackSpecificOptions(boolean isAudioEnabled, boolean isVideoEnabled) {
     audioMimeSpinner.setEnabled(isAudioEnabled);
@@ -291,6 +300,7 @@ public final class ConfigurationActivity extends AppCompatActivity {
     translateSpinner.setEnabled(isVideoEnabled);
     scaleSpinner.setEnabled(isVideoEnabled);
     rotateSpinner.setEnabled(isVideoEnabled);
+    enableHdrEditingCheckBox.setEnabled(isVideoEnabled);
 
     findViewById(R.id.audio_mime_text_view).setEnabled(isAudioEnabled);
     findViewById(R.id.video_mime_text_view).setEnabled(isVideoEnabled);
@@ -298,5 +308,6 @@ public final class ConfigurationActivity extends AppCompatActivity {
     findViewById(R.id.translate).setEnabled(isVideoEnabled);
     findViewById(R.id.scale).setEnabled(isVideoEnabled);
     findViewById(R.id.rotate).setEnabled(isVideoEnabled);
+    findViewById(R.id.hdr_editing).setEnabled(isVideoEnabled);
   }
 }
diff --git a/demos/transformer/src/main/java/com/google/android/exoplayer2/transformerdemo/TransformerActivity.java b/demos/transformer/src/main/java/com/google/android/exoplayer2/transformerdemo/TransformerActivity.java
index 888ddd3edc..8220bfe498 100644
--- a/demos/transformer/src/main/java/com/google/android/exoplayer2/transformerdemo/TransformerActivity.java
+++ b/demos/transformer/src/main/java/com/google/android/exoplayer2/transformerdemo/TransformerActivity.java
@@ -213,6 +213,8 @@ public final class TransformerActivity extends AppCompatActivity {
       if (!transformationMatrix.isIdentity()) {
         requestBuilder.setTransformationMatrix(transformationMatrix);
       }
+      requestBuilder.experimental_setEnableHdrEditing(
+          bundle.getBoolean(ConfigurationActivity.ENABLE_HDR_EDITING));
       transformerBuilder
           .setTransformationRequest(requestBuilder.build())
           .setRemoveAudio(bundle.getBoolean(ConfigurationActivity.SHOULD_REMOVE_AUDIO))
diff --git a/demos/transformer/src/main/res/layout/configuration_activity.xml b/demos/transformer/src/main/res/layout/configuration_activity.xml
index a9a9410a35..c3403a6269 100644
--- a/demos/transformer/src/main/res/layout/configuration_activity.xml
+++ b/demos/transformer/src/main/res/layout/configuration_activity.xml
@@ -164,6 +164,16 @@
         android:layout_gravity="right|center_vertical"
         android:gravity="right" />
     </TableRow>
+    <TableRow
+      android:layout_weight="1"
+      android:gravity="center_vertical" >
+      <TextView
+        android:id="@+id/hdr_editing"
+        android:text="@string/hdr_editing" />
+      <CheckBox
+        android:id="@+id/hdr_editing_checkbox"
+        android:layout_gravity="right" />
+    </TableRow>
   </TableLayout>
   <Button
     android:id="@+id/transform_button"
diff --git a/demos/transformer/src/main/res/values/strings.xml b/demos/transformer/src/main/res/values/strings.xml
index 39bd87ad8f..3c2c48c2f4 100644
--- a/demos/transformer/src/main/res/values/strings.xml
+++ b/demos/transformer/src/main/res/values/strings.xml
@@ -28,6 +28,7 @@
   <string name="scale" translatable="false">Scale video</string>
   <string name="rotate" translatable="false">Rotate video (degrees)</string>
   <string name="transform" translatable="false">Transform</string>
+  <string name="hdr_editing" translatable="false">[Experimental] HDR editing</string>
   <string name="debug_preview" translatable="false">Debug preview:</string>
   <string name="debug_preview_not_available" translatable="false">No debug preview available</string>
   <string name="transformation_started" translatable="false">Transformation started</string>
diff --git a/library/common/src/main/java/com/google/android/exoplayer2/util/GlUtil.java b/library/common/src/main/java/com/google/android/exoplayer2/util/GlUtil.java
index 6bdb8d17f2..487129bde3 100644
--- a/library/common/src/main/java/com/google/android/exoplayer2/util/GlUtil.java
+++ b/library/common/src/main/java/com/google/android/exoplayer2/util/GlUtil.java
@@ -42,7 +42,8 @@ import java.util.HashMap;
 import java.util.Map;
 import javax.microedition.khronos.egl.EGL10;
 
-/** OpenGL ES 2.0 utilities. */
+/** OpenGL ES utilities. */
+@SuppressWarnings("InlinedApi") // GLES constants are used safely based on the API version.
 public final class GlUtil {
 
   /** Thrown when an OpenGL error occurs and {@link #glAssertionsEnabled} is {@code true}. */
@@ -207,9 +208,44 @@ public final class GlUtil {
 
   private static final String TAG = "GlUtil";
 
+  // https://www.khronos.org/registry/EGL/extensions/EXT/EGL_EXT_protected_content.txt
   private static final String EXTENSION_PROTECTED_CONTENT = "EGL_EXT_protected_content";
+  // https://www.khronos.org/registry/EGL/extensions/KHR/EGL_KHR_surfaceless_context.txt
   private static final String EXTENSION_SURFACELESS_CONTEXT = "EGL_KHR_surfaceless_context";
 
+  // https://www.khronos.org/registry/OpenGL/extensions/EXT/EXT_YUV_target.txt
+  private static final int GL_SAMPLER_EXTERNAL_2D_Y2Y_EXT = 0x8BE7;
+  // https://www.khronos.org/registry/EGL/extensions/KHR/EGL_KHR_gl_colorspace.txt
+  private static final int EGL_GL_COLORSPACE_KHR = 0x309D;
+  // https://www.khronos.org/registry/EGL/extensions/EXT/EGL_EXT_gl_colorspace_bt2020_linear.txt
+  private static final int EGL_GL_COLORSPACE_BT2020_PQ_EXT = 0x3340;
+
+  private static final int[] EGL_WINDOW_SURFACE_ATTRIBUTES_NONE = new int[] {EGL14.EGL_NONE};
+  private static final int[] EGL_WINDOW_SURFACE_ATTRIBUTES_BT2020_PQ =
+      new int[] {EGL_GL_COLORSPACE_KHR, EGL_GL_COLORSPACE_BT2020_PQ_EXT, EGL14.EGL_NONE};
+  private static final int[] EGL_CONFIG_ATTRIBUTES_RGBA_8888 =
+      new int[] {
+        EGL14.EGL_RENDERABLE_TYPE, EGL14.EGL_OPENGL_ES2_BIT,
+        EGL14.EGL_RED_SIZE, /* redSize= */ 8,
+        EGL14.EGL_GREEN_SIZE, /* greenSize= */ 8,
+        EGL14.EGL_BLUE_SIZE, /* blueSize= */ 8,
+        EGL14.EGL_ALPHA_SIZE, /* alphaSize= */ 8,
+        EGL14.EGL_DEPTH_SIZE, /* depthSize= */ 0,
+        EGL14.EGL_STENCIL_SIZE, /* stencilSize= */ 0,
+        EGL14.EGL_NONE
+      };
+  private static final int[] EGL_CONFIG_ATTRIBUTES_RGBA_1010102 =
+      new int[] {
+        EGL14.EGL_RENDERABLE_TYPE, EGL14.EGL_OPENGL_ES2_BIT,
+        EGL14.EGL_RED_SIZE, /* redSize= */ 10,
+        EGL14.EGL_GREEN_SIZE, /* greenSize= */ 10,
+        EGL14.EGL_BLUE_SIZE, /* blueSize= */ 10,
+        EGL14.EGL_ALPHA_SIZE, /* alphaSize= */ 2,
+        EGL14.EGL_DEPTH_SIZE, /* depthSize= */ 0,
+        EGL14.EGL_STENCIL_SIZE, /* stencilSize= */ 0,
+        EGL14.EGL_NONE
+      };
+
   /** Class only contains static methods. */
   private GlUtil() {}
 
@@ -282,7 +318,16 @@ public final class GlUtil {
   /** Returns a new {@link EGLContext} for the specified {@link EGLDisplay}. */
   @RequiresApi(17)
   public static EGLContext createEglContext(EGLDisplay eglDisplay) {
-    return Api17.createEglContext(eglDisplay);
+    return Api17.createEglContext(eglDisplay, /* version= */ 2, EGL_CONFIG_ATTRIBUTES_RGBA_8888);
+  }
+
+  /**
+   * Returns a new {@link EGLContext} for the specified {@link EGLDisplay}, requesting ES 3 and an
+   * RGBA 1010102 config.
+   */
+  @RequiresApi(17)
+  public static EGLContext createEglContextEs3Rgba1010102(EGLDisplay eglDisplay) {
+    return Api17.createEglContext(eglDisplay, /* version= */ 3, EGL_CONFIG_ATTRIBUTES_RGBA_1010102);
   }
 
   /**
@@ -293,7 +338,24 @@ public final class GlUtil {
    */
   @RequiresApi(17)
   public static EGLSurface getEglSurface(EGLDisplay eglDisplay, Object surface) {
-    return Api17.getEglSurface(eglDisplay, surface);
+    return Api17.getEglSurface(
+        eglDisplay, surface, EGL_CONFIG_ATTRIBUTES_RGBA_8888, EGL_WINDOW_SURFACE_ATTRIBUTES_NONE);
+  }
+
+  /**
+   * Returns a new {@link EGLSurface} wrapping the specified {@code surface}, for HDR rendering with
+   * Rec. 2020 color primaries and using the PQ transfer function.
+   *
+   * @param eglDisplay The {@link EGLDisplay} to attach the surface to.
+   * @param surface The surface to wrap; must be a surface, surface texture or surface holder.
+   */
+  @RequiresApi(17)
+  public static EGLSurface getEglSurfaceBt2020Pq(EGLDisplay eglDisplay, Object surface) {
+    return Api17.getEglSurface(
+        eglDisplay,
+        surface,
+        EGL_CONFIG_ATTRIBUTES_RGBA_1010102,
+        EGL_WINDOW_SURFACE_ATTRIBUTES_BT2020_PQ);
   }
 
   /**
@@ -601,6 +663,13 @@ public final class GlUtil {
         return;
       }
 
+      if (type == GLES20.GL_FLOAT_MAT3) {
+        GLES20.glUniformMatrix3fv(
+            location, /* count= */ 1, /* transpose= */ false, value, /* offset= */ 0);
+        checkGlError();
+        return;
+      }
+
       if (type == GLES20.GL_FLOAT_MAT4) {
         GLES20.glUniformMatrix4fv(
             location, /* count= */ 1, /* transpose= */ false, value, /* offset= */ 0);
@@ -612,7 +681,7 @@ public final class GlUtil {
         throw new IllegalStateException("No call to setSamplerTexId() before bind.");
       }
       GLES20.glActiveTexture(GLES20.GL_TEXTURE0 + unit);
-      if (type == GLES11Ext.GL_SAMPLER_EXTERNAL_OES) {
+      if (type == GLES11Ext.GL_SAMPLER_EXTERNAL_OES || type == GL_SAMPLER_EXTERNAL_2D_Y2Y_EXT) {
         GLES20.glBindTexture(GLES11Ext.GL_TEXTURE_EXTERNAL_OES, texId);
       } else if (type == GLES20.GL_SAMPLER_2D) {
         GLES20.glBindTexture(GLES20.GL_TEXTURE_2D, texId);
@@ -651,12 +720,13 @@ public final class GlUtil {
     }
 
     @DoNotInline
-    public static EGLContext createEglContext(EGLDisplay eglDisplay) {
-      int[] contextAttributes = {EGL14.EGL_CONTEXT_CLIENT_VERSION, 2, EGL14.EGL_NONE};
+    public static EGLContext createEglContext(
+        EGLDisplay eglDisplay, int version, int[] configAttributes) {
+      int[] contextAttributes = {EGL14.EGL_CONTEXT_CLIENT_VERSION, version, EGL14.EGL_NONE};
       EGLContext eglContext =
           EGL14.eglCreateContext(
               eglDisplay,
-              getEglConfig(eglDisplay),
+              getEglConfig(eglDisplay, configAttributes),
               EGL14.EGL_NO_CONTEXT,
               contextAttributes,
               /* offset= */ 0);
@@ -664,19 +734,24 @@ public final class GlUtil {
         EGL14.eglTerminate(eglDisplay);
         throwGlException(
             "eglCreateContext() failed to create a valid context. The device may not support EGL"
-                + " version 2");
+                + " version "
+                + version);
       }
       checkGlError();
       return eglContext;
     }
 
     @DoNotInline
-    public static EGLSurface getEglSurface(EGLDisplay eglDisplay, Object surface) {
+    public static EGLSurface getEglSurface(
+        EGLDisplay eglDisplay,
+        Object surface,
+        int[] configAttributes,
+        int[] windowSurfaceAttributes) {
       return EGL14.eglCreateWindowSurface(
           eglDisplay,
-          getEglConfig(eglDisplay),
+          getEglConfig(eglDisplay, configAttributes),
           surface,
-          new int[] {EGL14.EGL_NONE},
+          windowSurfaceAttributes,
           /* offset= */ 0);
     }
 
@@ -717,22 +792,11 @@ public final class GlUtil {
     }
 
     @DoNotInline
-    private static EGLConfig getEglConfig(EGLDisplay eglDisplay) {
-      int[] defaultConfiguration =
-          new int[] {
-            EGL14.EGL_RENDERABLE_TYPE, EGL14.EGL_OPENGL_ES2_BIT,
-            EGL14.EGL_RED_SIZE, /* redSize= */ 8,
-            EGL14.EGL_GREEN_SIZE, /* greenSize= */ 8,
-            EGL14.EGL_BLUE_SIZE, /* blueSize= */ 8,
-            EGL14.EGL_ALPHA_SIZE, /* alphaSize= */ 8,
-            EGL14.EGL_DEPTH_SIZE, /* depthSize= */ 0,
-            EGL14.EGL_STENCIL_SIZE, /* stencilSize= */ 0,
-            EGL14.EGL_NONE
-          };
+    private static EGLConfig getEglConfig(EGLDisplay eglDisplay, int[] attributes) {
       EGLConfig[] eglConfigs = new EGLConfig[1];
       if (!EGL14.eglChooseConfig(
           eglDisplay,
-          defaultConfiguration,
+          attributes,
           /* attrib_listOffset= */ 0,
           eglConfigs,
           /* configsOffset= */ 0,
diff --git a/library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/FrameEditorDataProcessingTest.java b/library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/FrameEditorDataProcessingTest.java
index 38471fa1d0..d03e731f4a 100644
--- a/library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/FrameEditorDataProcessingTest.java
+++ b/library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/FrameEditorDataProcessingTest.java
@@ -205,6 +205,7 @@ public final class FrameEditorDataProcessingTest {
               PIXEL_WIDTH_HEIGHT_RATIO,
               transformationMatrix,
               frameEditorOutputImageReader.getSurface(),
+              /* enableExperimentalHdrEditing= */ false,
               Transformer.DebugViewProvider.NONE);
       frameEditor.registerInputFrame();
 
diff --git a/library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/FrameEditorTest.java b/library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/FrameEditorTest.java
index a30ffc51c2..f6791f82ca 100644
--- a/library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/FrameEditorTest.java
+++ b/library/transformer/src/androidTest/java/com/google/android/exoplayer2/transformer/FrameEditorTest.java
@@ -28,7 +28,7 @@ import org.junit.Test;
 import org.junit.runner.RunWith;
 
 /**
- * Test for {@link FrameEditor#create(Context, int, int, float, Matrix, Surface,
+ * Test for {@link FrameEditor#create(Context, int, int, float, Matrix, Surface, boolean,
  * Transformer.DebugViewProvider) creating} a {@link FrameEditor}.
  */
 @RunWith(AndroidJUnit4.class)
@@ -46,6 +46,7 @@ public final class FrameEditorTest {
         /* pixelWidthHeightRatio= */ 1,
         new Matrix(),
         new Surface(new SurfaceTexture(false)),
+        /* enableExperimentalHdrEditing= */ false,
         Transformer.DebugViewProvider.NONE);
   }
 
@@ -62,6 +63,7 @@ public final class FrameEditorTest {
                     /* pixelWidthHeightRatio= */ 2,
                     new Matrix(),
                     new Surface(new SurfaceTexture(false)),
+                    /* enableExperimentalHdrEditing= */ false,
                     Transformer.DebugViewProvider.NONE));
 
     assertThat(exception).hasCauseThat().isInstanceOf(UnsupportedOperationException.class);
diff --git a/library/transformer/src/main/assets/shaders/fragment_shader.glsl b/library/transformer/src/main/assets/shaders/fragment_shader_copy_external.glsl
similarity index 100%
rename from library/transformer/src/main/assets/shaders/fragment_shader.glsl
rename to library/transformer/src/main/assets/shaders/fragment_shader_copy_external.glsl
diff --git a/library/transformer/src/main/assets/shaders/fragment_shader_copy_external_yuv_es3.glsl b/library/transformer/src/main/assets/shaders/fragment_shader_copy_external_yuv_es3.glsl
new file mode 100644
index 0000000000..1a6ea2b312
--- /dev/null
+++ b/library/transformer/src/main/assets/shaders/fragment_shader_copy_external_yuv_es3.glsl
@@ -0,0 +1,29 @@
+#version 300 es
+// Copyright 2022 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.
+#extension GL_OES_EGL_image_external : require
+#extension GL_EXT_YUV_target : require
+precision mediump float;
+uniform __samplerExternal2DY2YEXT uTexSampler;
+uniform mat3 uColorTransform;
+in vec2 vTexCoords;
+out vec4 outColor;
+void main() {
+  vec3 srcYuv = texture(uTexSampler, vTexCoords).xyz;
+  vec3 yuvOffset;
+  yuvOffset.x = srcYuv.r - 0.0625;
+  yuvOffset.y = srcYuv.g - 0.5;
+  yuvOffset.z = srcYuv.b - 0.5;
+  outColor = vec4(uColorTransform * yuvOffset, 1.0);
+}
diff --git a/library/transformer/src/main/assets/shaders/vertex_shader.glsl b/library/transformer/src/main/assets/shaders/vertex_shader_transformation.glsl
similarity index 100%
rename from library/transformer/src/main/assets/shaders/vertex_shader.glsl
rename to library/transformer/src/main/assets/shaders/vertex_shader_transformation.glsl
diff --git a/library/transformer/src/main/assets/shaders/vertex_shader_transformation_es3.glsl b/library/transformer/src/main/assets/shaders/vertex_shader_transformation_es3.glsl
new file mode 100644
index 0000000000..9af4dcecd6
--- /dev/null
+++ b/library/transformer/src/main/assets/shaders/vertex_shader_transformation_es3.glsl
@@ -0,0 +1,23 @@
+#version 300 es
+// Copyright 2022 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.
+in vec4 aFramePosition;
+in vec4 aTexCoords;
+uniform mat4 uTexTransform;
+uniform mat4 uTransformationMatrix;
+out vec2 vTexCoords;
+void main() {
+  gl_Position = uTransformationMatrix * aFramePosition;
+  vTexCoords = (uTexTransform * aTexCoords).xy;
+}
diff --git a/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/FrameEditor.java b/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/FrameEditor.java
index 4a2fcf39c0..fea7a522f1 100644
--- a/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/FrameEditor.java
+++ b/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/FrameEditor.java
@@ -51,6 +51,7 @@ import java.util.concurrent.atomic.AtomicInteger;
    * @param pixelWidthHeightRatio The ratio of width over height, for each pixel.
    * @param transformationMatrix The transformation matrix to apply to each frame.
    * @param outputSurface The {@link Surface}.
+   * @param enableExperimentalHdrEditing Whether to attempt to process the input as an HDR signal.
    * @param debugViewProvider Provider for optional debug views to show intermediate output.
    * @return A configured {@code FrameEditor}.
    * @throws TransformationException If the {@code pixelWidthHeightRatio} isn't 1, reading shader
@@ -64,6 +65,7 @@ import java.util.concurrent.atomic.AtomicInteger;
       float pixelWidthHeightRatio,
       Matrix transformationMatrix,
       Surface outputSurface,
+      boolean enableExperimentalHdrEditing,
       Transformer.DebugViewProvider debugViewProvider)
       throws TransformationException {
     if (pixelWidthHeightRatio != 1.0f) {
@@ -85,18 +87,32 @@ import java.util.concurrent.atomic.AtomicInteger;
     EGLSurface eglSurface;
     int textureId;
     GlUtil.Program glProgram;
-    @Nullable EGLSurface debugPreviewEglSurface;
+    @Nullable EGLSurface debugPreviewEglSurface = null;
     try {
       eglDisplay = GlUtil.createEglDisplay();
-      eglContext = GlUtil.createEglContext(eglDisplay);
-      eglSurface = GlUtil.getEglSurface(eglDisplay, outputSurface);
+
+      if (enableExperimentalHdrEditing) {
+        eglContext = GlUtil.createEglContextEs3Rgba1010102(eglDisplay);
+        // TODO(b/209404935): Don't assume BT.2020 PQ input/output.
+        eglSurface = GlUtil.getEglSurfaceBt2020Pq(eglDisplay, outputSurface);
+        if (debugSurfaceView != null) {
+          debugPreviewEglSurface =
+              GlUtil.getEglSurfaceBt2020Pq(eglDisplay, checkNotNull(debugSurfaceView.getHolder()));
+        }
+      } else {
+        eglContext = GlUtil.createEglContext(eglDisplay);
+        eglSurface = GlUtil.getEglSurface(eglDisplay, outputSurface);
+        if (debugSurfaceView != null) {
+          debugPreviewEglSurface =
+              GlUtil.getEglSurface(eglDisplay, checkNotNull(debugSurfaceView.getHolder()));
+        }
+      }
+
       GlUtil.focusSurface(eglDisplay, eglContext, eglSurface, outputWidth, outputHeight);
       textureId = GlUtil.createExternalTexture();
-      glProgram = configureGlProgram(context, transformationMatrix, textureId);
-      debugPreviewEglSurface =
-          debugSurfaceView == null
-              ? null
-              : GlUtil.getEglSurface(eglDisplay, checkNotNull(debugSurfaceView.getHolder()));
+      glProgram =
+          configureGlProgram(
+              context, transformationMatrix, textureId, enableExperimentalHdrEditing);
     } catch (IOException | GlUtil.GlException e) {
       throw TransformationException.createForFrameEditor(
           e, TransformationException.ERROR_CODE_GL_INIT_FAILED);
@@ -126,11 +142,23 @@ import java.util.concurrent.atomic.AtomicInteger;
   }
 
   private static GlUtil.Program configureGlProgram(
-      Context context, Matrix transformationMatrix, int textureId) throws IOException {
+      Context context,
+      Matrix transformationMatrix,
+      int textureId,
+      boolean enableExperimentalHdrEditing)
+      throws IOException {
     // TODO(b/205002913): check the loaded program is consistent with the attributes
     // and uniforms expected in the code.
+    String vertexShaderFilePath =
+        enableExperimentalHdrEditing
+            ? VERTEX_SHADER_TRANSFORMATION_ES3_PATH
+            : VERTEX_SHADER_TRANSFORMATION_PATH;
+    String fragmentShaderFilePath =
+        enableExperimentalHdrEditing
+            ? FRAGMENT_SHADER_COPY_EXTERNAL_YUV_ES3_PATH
+            : FRAGMENT_SHADER_COPY_EXTERNAL_PATH;
     GlUtil.Program glProgram =
-        new GlUtil.Program(context, VERTEX_SHADER_FILE_PATH, FRAGMENT_SHADER_FILE_PATH);
+        new GlUtil.Program(context, vertexShaderFilePath, fragmentShaderFilePath);
 
     // Draw the frame on the entire normalized device coordinate space, from -1 to 1, for x and y.
     glProgram.setBufferAttribute(
@@ -139,6 +167,11 @@ import java.util.concurrent.atomic.AtomicInteger;
         "aTexCoords", GlUtil.getTextureCoordinateBounds(), GlUtil.RECTANGLE_VERTICES_COUNT);
     glProgram.setSamplerTexIdUniform("uTexSampler", textureId, /* unit= */ 0);
 
+    if (enableExperimentalHdrEditing) {
+      // In HDR editing mode the decoder output is sampled in YUV.
+      glProgram.setFloatsUniform("uColorTransform", MATRIX_YUV_TO_BT2020_COLOR_TRANSFORM);
+    }
+
     float[] transformationMatrixArray = getGlMatrixArray(transformationMatrix);
     glProgram.setFloatsUniform("uTransformationMatrix", transformationMatrixArray);
     return glProgram;
@@ -184,9 +217,21 @@ import java.util.concurrent.atomic.AtomicInteger;
     return matrix4x4Array;
   }
 
-  // Predefined shader values.
-  private static final String VERTEX_SHADER_FILE_PATH = "shaders/vertex_shader.glsl";
-  private static final String FRAGMENT_SHADER_FILE_PATH = "shaders/fragment_shader.glsl";
+  private static final String VERTEX_SHADER_TRANSFORMATION_PATH =
+      "shaders/vertex_shader_transformation.glsl";
+  private static final String FRAGMENT_SHADER_COPY_EXTERNAL_PATH =
+      "shaders/fragment_shader_copy_external.glsl";
+  private static final String VERTEX_SHADER_TRANSFORMATION_ES3_PATH =
+      "shaders/vertex_shader_transformation_es3.glsl";
+  private static final String FRAGMENT_SHADER_COPY_EXTERNAL_YUV_ES3_PATH =
+      "shaders/fragment_shader_copy_external_yuv_es3.glsl";
+  // Color transform coefficients from
+  // https://cs.android.com/android/platform/superproject/+/master:frameworks/av/media/libstagefright/colorconversion/ColorConverter.cpp;l=668-670;drc=487adf977a50cac3929eba15fad0d0f461c7ff0f.
+  private static final float[] MATRIX_YUV_TO_BT2020_COLOR_TRANSFORM = {
+    1.168f, 1.168f, 1.168f,
+    0.0f, -0.188f, 2.148f,
+    1.683f, -0.652f, 0.0f,
+  };
 
   private final float[] textureTransformMatrix;
   private final EGLDisplay eglDisplay;
diff --git a/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/TransformationRequest.java b/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/TransformationRequest.java
index d756dbffe4..20f249c359 100644
--- a/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/TransformationRequest.java
+++ b/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/TransformationRequest.java
@@ -41,6 +41,7 @@ public final class TransformationRequest {
     private int outputHeight;
     @Nullable private String audioMimeType;
     @Nullable private String videoMimeType;
+    private boolean enableHdrEditing;
 
     /**
      * Creates a new instance with default values.
@@ -59,6 +60,7 @@ public final class TransformationRequest {
       this.outputHeight = transformationRequest.outputHeight;
       this.audioMimeType = transformationRequest.audioMimeType;
       this.videoMimeType = transformationRequest.videoMimeType;
+      this.enableHdrEditing = transformationRequest.enableHdrEditing;
     }
 
     /**
@@ -192,10 +194,32 @@ public final class TransformationRequest {
       return this;
     }
 
+    /**
+     * Sets whether to attempt to process any input video stream as a high dynamic range (HDR)
+     * signal.
+     *
+     * <p>This method is experimental, and will be renamed or removed in a future release. The HDR
+     * editing feature is under development and is intended for developing/testing HDR processing
+     * and encoding support.
+     *
+     * @param enableHdrEditing Whether to attempt to process any input video stream as a high
+     *     dynamic range (HDR) signal.
+     * @return This builder.
+     */
+    public Builder experimental_setEnableHdrEditing(boolean enableHdrEditing) {
+      this.enableHdrEditing = enableHdrEditing;
+      return this;
+    }
+
     /** Builds a {@link TransformationRequest} instance. */
     public TransformationRequest build() {
       return new TransformationRequest(
-          transformationMatrix, flattenForSlowMotion, outputHeight, audioMimeType, videoMimeType);
+          transformationMatrix,
+          flattenForSlowMotion,
+          outputHeight,
+          audioMimeType,
+          videoMimeType,
+          enableHdrEditing);
     }
   }
 
@@ -231,18 +255,26 @@ public final class TransformationRequest {
    * @see Builder#setVideoMimeType(String)
    */
   @Nullable public final String videoMimeType;
+  /**
+   * Whether to attempt to process any input video stream as a high dynamic range (HDR) signal.
+   *
+   * @see Builder#experimental_setEnableHdrEditing(boolean)
+   */
+  public final boolean enableHdrEditing;
 
   private TransformationRequest(
       Matrix transformationMatrix,
       boolean flattenForSlowMotion,
       int outputHeight,
       @Nullable String audioMimeType,
-      @Nullable String videoMimeType) {
+      @Nullable String videoMimeType,
+      boolean enableHdrEditing) {
     this.transformationMatrix = transformationMatrix;
     this.flattenForSlowMotion = flattenForSlowMotion;
     this.outputHeight = outputHeight;
     this.audioMimeType = audioMimeType;
     this.videoMimeType = videoMimeType;
+    this.enableHdrEditing = enableHdrEditing;
   }
 
   @Override
@@ -258,7 +290,8 @@ public final class TransformationRequest {
         && flattenForSlowMotion == that.flattenForSlowMotion
         && outputHeight == that.outputHeight
         && Util.areEqual(audioMimeType, that.audioMimeType)
-        && Util.areEqual(videoMimeType, that.videoMimeType);
+        && Util.areEqual(videoMimeType, that.videoMimeType)
+        && enableHdrEditing == that.enableHdrEditing;
   }
 
   @Override
@@ -268,6 +301,7 @@ public final class TransformationRequest {
     result = 31 * result + outputHeight;
     result = 31 * result + (audioMimeType != null ? audioMimeType.hashCode() : 0);
     result = 31 * result + (videoMimeType != null ? videoMimeType.hashCode() : 0);
+    result = 31 * result + (enableHdrEditing ? 1 : 0);
     return result;
   }
 
diff --git a/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/TransformerVideoRenderer.java b/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/TransformerVideoRenderer.java
index 062fecb0ab..8299fda5a5 100644
--- a/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/TransformerVideoRenderer.java
+++ b/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/TransformerVideoRenderer.java
@@ -99,6 +99,9 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
   }
 
   private boolean shouldPassthrough(Format inputFormat) {
+    if (transformationRequest.enableHdrEditing) {
+      return false;
+    }
     if (transformationRequest.videoMimeType != null
         && !transformationRequest.videoMimeType.equals(inputFormat.sampleMimeType)) {
       return false;
diff --git a/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/VideoTranscodingSamplePipeline.java b/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/VideoTranscodingSamplePipeline.java
index 45d04d7d66..1613b06e69 100644
--- a/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/VideoTranscodingSamplePipeline.java
+++ b/library/transformer/src/main/java/com/google/android/exoplayer2/transformer/VideoTranscodingSamplePipeline.java
@@ -125,7 +125,8 @@ import org.checkerframework.dataflow.qual.Pure;
             requestedOutputFormat,
             actualOutputFormat));
 
-    if (inputFormat.height != actualOutputFormat.height
+    if (transformationRequest.enableHdrEditing
+        || inputFormat.height != actualOutputFormat.height
         || inputFormat.width != actualOutputFormat.width
         || !transformationMatrix.isIdentity()) {
       frameEditor =
@@ -136,6 +137,7 @@ import org.checkerframework.dataflow.qual.Pure;
               inputFormat.pixelWidthHeightRatio,
               transformationMatrix,
               /* outputSurface= */ checkNotNull(encoder.getInputSurface()),
+              transformationRequest.enableHdrEditing,
               debugViewProvider);
     }