From 655ca9f343181b29739e86dd169eab393c898dd8 Mon Sep 17 00:00:00 2001
From: Gopalakrishnan Nallasamy <gokrishnan@google.com>
Date: Tue, 16 Nov 2021 19:04:07 -0800
Subject: [PATCH] MediaCodec:Add code for YCbCrP010 format support

Set MediaFormat based on bitDepth and bitDepthAllocated.
Validate semiplanarness of chroma planes.

Bug: 206456921

Test: atest android.media.cts.ImageReaderDecoderTest
Change-Id: I88145d7eb9d372deb02ca2d633eebe087f7c4733
---
 media/java/android/media/Image.java      |  8 +++++
 media/java/android/media/MediaCodec.java | 42 ++++++++++++++++++++++--
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/media/java/android/media/Image.java b/media/java/android/media/Image.java
index 1616c03112a8f..54e19db8e7e92 100644
--- a/media/java/android/media/Image.java
+++ b/media/java/android/media/Image.java
@@ -163,6 +163,14 @@ public abstract class Image implements AutoCloseable {
      *      {@link android.graphics.BitmapFactory#decodeByteArray BitmapFactory#decodeByteArray}.
      *   </td>
      * </tr>
+     * <tr>
+     *   <td>{@link android.graphics.ImageFormat#YCBCR_P010 YCBCR_P010}</td>
+     *   <td>1</td>
+     *   <td>P010 is a 4:2:0 YCbCr semiplanar format comprised of a WxH Y plane
+     *     followed by a Wx(H/2) CbCr plane. Each sample is represented by a 16-bit
+     *     little-endian value, with the lower 6 bits set to zero.
+     *   </td>
+     * </tr>
      * </table>
      *
      * @see android.graphics.ImageFormat
diff --git a/media/java/android/media/MediaCodec.java b/media/java/android/media/MediaCodec.java
index de31a7f5c1842..72dd2bd43eb62 100644
--- a/media/java/android/media/MediaCodec.java
+++ b/media/java/android/media/MediaCodec.java
@@ -5107,7 +5107,6 @@ final public class MediaCodec {
         public MediaImage(
                 @NonNull ByteBuffer buffer, @NonNull ByteBuffer info, boolean readOnly,
                 long timestamp, int xOffset, int yOffset, @Nullable Rect cropRect) {
-            mFormat = ImageFormat.YUV_420_888;
             mTimestamp = timestamp;
             mIsImageValid = true;
             mIsReadOnly = buffer.isReadOnly();
@@ -5120,6 +5119,11 @@ final public class MediaCodec {
 
             mBufferContext = 0;
 
+            int cbPlaneOffset = -1;
+            int crPlaneOffset = -1;
+            int planeOffsetInc = -1;
+            int pixelStride = -1;
+
             // read media-info.  See MediaImage2
             if (info.remaining() == 104) {
                 int type = info.getInt();
@@ -5137,14 +5141,27 @@ final public class MediaCodec {
                             "unsupported size: " + mWidth + "x" + mHeight);
                 }
                 int bitDepth = info.getInt();
-                if (bitDepth != 8) {
+                if (bitDepth != 8 && bitDepth != 10) {
                     throw new UnsupportedOperationException("unsupported bit depth: " + bitDepth);
                 }
                 int bitDepthAllocated = info.getInt();
-                if (bitDepthAllocated != 8) {
+                if (bitDepthAllocated != 8 && bitDepthAllocated != 16) {
                     throw new UnsupportedOperationException(
                             "unsupported allocated bit depth: " + bitDepthAllocated);
                 }
+                if (bitDepth == 8 && bitDepthAllocated == 8) {
+                    mFormat = ImageFormat.YUV_420_888;
+                    planeOffsetInc = 1;
+                    pixelStride = 2;
+                } else if (bitDepth == 10 && bitDepthAllocated == 16) {
+                    mFormat = ImageFormat.YCBCR_P010;
+                    planeOffsetInc = 2;
+                    pixelStride = 4;
+                } else {
+                    throw new UnsupportedOperationException("couldn't infer ImageFormat"
+                      + " bitDepth: " + bitDepth + " bitDepthAllocated: " + bitDepthAllocated);
+                }
+
                 mPlanes = new MediaPlane[numPlanes];
                 for (int ix = 0; ix < numPlanes; ix++) {
                     int planeOffset = info.getInt();
@@ -5166,12 +5183,31 @@ final public class MediaCodec {
                     buffer.limit(buffer.position() + Utils.divUp(bitDepth, 8)
                             + (mHeight / vert - 1) * rowInc + (mWidth / horiz - 1) * colInc);
                     mPlanes[ix] = new MediaPlane(buffer.slice(), rowInc, colInc);
+                    if ((mFormat == ImageFormat.YUV_420_888 || mFormat == ImageFormat.YCBCR_P010)
+                            && ix == 1) {
+                        cbPlaneOffset = planeOffset;
+                    } else if ((mFormat == ImageFormat.YUV_420_888
+                            || mFormat == ImageFormat.YCBCR_P010) && ix == 2) {
+                        crPlaneOffset = planeOffset;
+                    }
                 }
             } else {
                 throw new UnsupportedOperationException(
                         "unsupported info length: " + info.remaining());
             }
 
+            // Validate chroma semiplanerness.
+            if (mFormat == ImageFormat.YCBCR_P010) {
+                if (crPlaneOffset != cbPlaneOffset + planeOffsetInc) {
+                    throw new UnsupportedOperationException("Invalid plane offsets"
+                    + " cbPlaneOffset: " + cbPlaneOffset + " crPlaneOffset: " + crPlaneOffset);
+                }
+                if (mPlanes[1].getPixelStride() != pixelStride
+                        || mPlanes[2].getPixelStride() != pixelStride) {
+                    throw new UnsupportedOperationException("Invalid pixelStride");
+                }
+            }
+
             if (cropRect == null) {
                 cropRect = new Rect(0, 0, mWidth, mHeight);
             }