JPEG/R: stride support

Bug: 264715926
Test: jpegr_test.cpp, manual test with camera/display app
Change-Id: I92c1d19b8bd49060794b4a9427c67017aa6743a9
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegr.h b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegr.h
index 1ab1dd7..766d473 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegr.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegr.h
@@ -59,12 +59,26 @@
 struct jpegr_uncompressed_struct {
     // Pointer to the data location.
     void* data;
-    // Width of the recovery map or image in pixels.
+    // Width of the recovery map or the luma plane of the image in pixels.
     int width;
-    // Height of the recovery map or image in pixels.
+    // Height of the recovery map or the luma plane of the image in pixels.
     int height;
     // Color gamut.
     jpegr_color_gamut colorGamut;
+
+    // Values below are optional
+    // Pointer to chroma data, if it's NULL, chroma plane is considered to be immediately
+    // following after the luma plane.
+    // Note: currently this feature is only supported for P010 image (HDR input).
+    void* chroma_data = nullptr;
+    // Strides of Y plane in number of pixels, using 0 to present uninitialized, must be
+    // larger than or equal to luma width.
+    // Note: currently this feature is only supported for P010 image (HDR input).
+    int luma_stride = 0;
+    // Strides of UV plane in number of pixels, using 0 to present uninitialized, must be
+    // larger than or equal to chroma width.
+    // Note: currently this feature is only supported for P010 image (HDR input).
+    int chroma_stride = 0;
 };
 
 /*
@@ -358,6 +372,16 @@
      */
     status_t toneMap(jr_uncompressed_ptr src,
                      jr_uncompressed_ptr dest);
+
+    /*
+     * This method will check the validity of the input images.
+     *
+     * @param uncompressed_p010_image uncompressed HDR image in P010 color format
+     * @param uncompressed_yuv_420_image uncompressed SDR image in YUV_420 color format
+     * @return NO_ERROR if the input images are valid, error code is not valid.
+     */
+    status_t areInputImagesValid(jr_uncompressed_ptr uncompressed_p010_image,
+                                 jr_uncompressed_ptr uncompressed_yuv_420_image);
 };
 
 } // namespace android::jpegrecoverymap
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegrerrorcode.h b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegrerrorcode.h
index f730343..159aaa8 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegrerrorcode.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegrerrorcode.h
@@ -46,6 +46,8 @@
     ERROR_JPEGR_CALCULATION_ERROR       = JPEGR_RUNTIME_ERROR_BASE - 3,
     ERROR_JPEGR_METADATA_ERROR          = JPEGR_RUNTIME_ERROR_BASE - 4,
     ERROR_JPEGR_TONEMAP_ERROR           = JPEGR_RUNTIME_ERROR_BASE - 5,
+
+    ERROR_JPEGR_UNSUPPORTED_FEATURE     = -20000,
 };
 
 }  // namespace android::jpegrecoverymap
diff --git a/libs/jpegrecoverymap/jpegr.cpp b/libs/jpegrecoverymap/jpegr.cpp
index e395d51..646b79f 100644
--- a/libs/jpegrecoverymap/jpegr.cpp
+++ b/libs/jpegrecoverymap/jpegr.cpp
@@ -86,6 +86,54 @@
   return cpuCoreCount;
 }
 
+status_t JpegR::areInputImagesValid(jr_uncompressed_ptr uncompressed_p010_image,
+                                    jr_uncompressed_ptr uncompressed_yuv_420_image) {
+  if (uncompressed_p010_image == nullptr) {
+    return ERROR_JPEGR_INVALID_NULL_PTR;
+  }
+
+  if (uncompressed_p010_image->width % kJpegBlock != 0
+          || uncompressed_p010_image->height % 2 != 0) {
+    ALOGE("Image size can not be handled: %dx%d.",
+            uncompressed_p010_image->width, uncompressed_p010_image->height);
+    return ERROR_JPEGR_INVALID_INPUT_TYPE;
+  }
+
+  if (uncompressed_p010_image->luma_stride != 0
+          && uncompressed_p010_image->luma_stride < uncompressed_p010_image->width) {
+    ALOGE("Image stride can not be smaller than width, stride=%d, width=%d",
+                uncompressed_p010_image->luma_stride, uncompressed_p010_image->width);
+    return ERROR_JPEGR_INVALID_INPUT_TYPE;
+  }
+
+  if (uncompressed_yuv_420_image == nullptr) {
+    return NO_ERROR;
+  }
+
+  if (uncompressed_yuv_420_image->luma_stride != 0) {
+    ALOGE("Stride is not supported for YUV420 image");
+    return ERROR_JPEGR_UNSUPPORTED_FEATURE;
+  }
+
+  if (uncompressed_yuv_420_image->chroma_data != nullptr) {
+    ALOGE("Pointer to chroma plane is not supported for YUV420 image, chroma data must"
+          "be immediately after the luma data.");
+    return ERROR_JPEGR_UNSUPPORTED_FEATURE;
+  }
+
+  if (uncompressed_p010_image->width != uncompressed_yuv_420_image->width
+      || uncompressed_p010_image->height != uncompressed_yuv_420_image->height) {
+    ALOGE("Image resolutions mismatch: P010: %dx%d, YUV420: %dx%d",
+              uncompressed_p010_image->width,
+              uncompressed_p010_image->height,
+              uncompressed_yuv_420_image->width,
+              uncompressed_yuv_420_image->height);
+    return ERROR_JPEGR_RESOLUTION_MISMATCH;
+  }
+
+  return NO_ERROR;
+}
+
 /* Encode API-0 */
 status_t JpegR::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
                             jpegr_transfer_function hdr_tf,
@@ -100,11 +148,9 @@
     return ERROR_JPEGR_INVALID_INPUT_TYPE;
   }
 
-  if (uncompressed_p010_image->width % kJpegBlock != 0
-          || uncompressed_p010_image->height % 2 != 0) {
-    ALOGE("Image size can not be handled: %dx%d",
-            uncompressed_p010_image->width, uncompressed_p010_image->height);
-    return ERROR_JPEGR_INVALID_INPUT_TYPE;
+  if (status_t ret = areInputImagesValid(
+          uncompressed_p010_image, /* uncompressed_yuv_420_image */ nullptr) != NO_ERROR) {
+    return ret;
   }
 
   jpegr_metadata_struct metadata;
@@ -164,16 +210,9 @@
     return ERROR_JPEGR_INVALID_INPUT_TYPE;
   }
 
-  if (uncompressed_p010_image->width != uncompressed_yuv_420_image->width
-   || uncompressed_p010_image->height != uncompressed_yuv_420_image->height) {
-    return ERROR_JPEGR_RESOLUTION_MISMATCH;
-  }
-
-  if (uncompressed_p010_image->width % kJpegBlock != 0
-          || uncompressed_p010_image->height % 2 != 0) {
-    ALOGE("Image size can not be handled: %dx%d",
-            uncompressed_p010_image->width, uncompressed_p010_image->height);
-    return ERROR_JPEGR_INVALID_INPUT_TYPE;
+  if (status_t ret = areInputImagesValid(
+          uncompressed_p010_image, uncompressed_yuv_420_image) != NO_ERROR) {
+    return ret;
   }
 
   jpegr_metadata_struct metadata;
@@ -223,16 +262,9 @@
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  if (uncompressed_p010_image->width != uncompressed_yuv_420_image->width
-   || uncompressed_p010_image->height != uncompressed_yuv_420_image->height) {
-    return ERROR_JPEGR_RESOLUTION_MISMATCH;
-  }
-
-  if (uncompressed_p010_image->width % kJpegBlock != 0
-          || uncompressed_p010_image->height % 2 != 0) {
-    ALOGE("Image size can not be handled: %dx%d",
-            uncompressed_p010_image->width, uncompressed_p010_image->height);
-    return ERROR_JPEGR_INVALID_INPUT_TYPE;
+  if (status_t ret = areInputImagesValid(
+          uncompressed_p010_image, uncompressed_yuv_420_image) != NO_ERROR) {
+    return ret;
   }
 
   jpegr_metadata_struct metadata;
@@ -266,11 +298,9 @@
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  if (uncompressed_p010_image->width % kJpegBlock != 0
-          || uncompressed_p010_image->height % 2 != 0) {
-    ALOGE("Image size can not be handled: %dx%d",
-            uncompressed_p010_image->width, uncompressed_p010_image->height);
-    return ERROR_JPEGR_INVALID_INPUT_TYPE;
+  if (status_t ret = areInputImagesValid(
+          uncompressed_p010_image, /* uncompressed_yuv_420_image */ nullptr) != NO_ERROR) {
+    return ret;
   }
 
   JpegDecoderHelper jpeg_decoder;
@@ -993,25 +1023,43 @@
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
+  size_t src_luma_stride = src->luma_stride;
+  size_t src_chroma_stride = src->chroma_stride;
+  uint16_t* src_luma_data = reinterpret_cast<uint16_t*>(src->data);
+  uint16_t* src_chroma_data = reinterpret_cast<uint16_t*>(src->chroma_data);
+
+  if (src_chroma_data == nullptr) {
+    src_chroma_data = &reinterpret_cast<uint16_t*>(src->data)[src->luma_stride * src->height];
+  }
+  if (src_luma_stride == 0) {
+    src_luma_stride = src->width;
+  }
+  if (src_chroma_stride == 0) {
+    src_chroma_stride = src_luma_stride;
+  }
+
   dest->width = src->width;
   dest->height = src->height;
 
-  size_t pixel_count = src->width * src->height;
+  size_t dest_luma_pixel_count = dest->width * dest->height;
+
   for (size_t y = 0; y < src->height; ++y) {
     for (size_t x = 0; x < src->width; ++x) {
-      size_t pixel_y_idx = x + y * src->width;
-      size_t pixel_uv_idx = x / 2 + (y / 2) * (src->width / 2);
+      size_t src_y_idx = y * src_luma_stride + x;
+      size_t src_u_idx = (y >> 1) * src_chroma_stride + (x & ~0x1);
+      size_t src_v_idx = src_u_idx + 1;
 
-      uint16_t y_uint = reinterpret_cast<uint16_t*>(src->data)[pixel_y_idx]
-                        >> 6;
-      uint16_t u_uint = reinterpret_cast<uint16_t*>(src->data)[pixel_count + pixel_uv_idx * 2]
-                        >> 6;
-      uint16_t v_uint = reinterpret_cast<uint16_t*>(src->data)[pixel_count + pixel_uv_idx * 2 + 1]
-                        >> 6;
+      uint16_t y_uint = src_luma_data[src_y_idx] >> 6;
+      uint16_t u_uint = src_chroma_data[src_u_idx] >> 6;
+      uint16_t v_uint = src_chroma_data[src_v_idx] >> 6;
 
-      uint8_t* y = &reinterpret_cast<uint8_t*>(dest->data)[pixel_y_idx];
-      uint8_t* u = &reinterpret_cast<uint8_t*>(dest->data)[pixel_count + pixel_uv_idx];
-      uint8_t* v = &reinterpret_cast<uint8_t*>(dest->data)[pixel_count * 5 / 4 + pixel_uv_idx];
+      size_t dest_y_idx = x + y * dest->width;
+      size_t dest_uv_idx = x / 2 + (y / 2) * (dest->width / 2);
+
+      uint8_t* y = &reinterpret_cast<uint8_t*>(dest->data)[dest_y_idx];
+      uint8_t* u = &reinterpret_cast<uint8_t*>(dest->data)[dest_luma_pixel_count + dest_uv_idx];
+      uint8_t* v = &reinterpret_cast<uint8_t*>(
+              dest->data)[dest_luma_pixel_count * 5 / 4 + dest_uv_idx];
 
       *y = static_cast<uint8_t>((y_uint >> 2) & 0xff);
       *u = static_cast<uint8_t>((u_uint >> 2) & 0xff);
diff --git a/libs/jpegrecoverymap/recoverymapmath.cpp b/libs/jpegrecoverymap/recoverymapmath.cpp
index 2cffde3..7c0c15a 100644
--- a/libs/jpegrecoverymap/recoverymapmath.cpp
+++ b/libs/jpegrecoverymap/recoverymapmath.cpp
@@ -493,17 +493,28 @@
 }
 
 Color getP010Pixel(jr_uncompressed_ptr image, size_t x, size_t y) {
-  size_t pixel_count = image->width * image->height;
+  size_t luma_stride = image->luma_stride;
+  size_t chroma_stride = image->chroma_stride;
+  uint16_t* luma_data = reinterpret_cast<uint16_t*>(image->data);
+  uint16_t* chroma_data = reinterpret_cast<uint16_t*>(image->chroma_data);
 
-  size_t pixel_y_idx = x + y * image->width;
-  size_t pixel_uv_idx = x / 2 + (y / 2) * (image->width / 2);
+  if (luma_stride == 0) {
+    luma_stride = image->width;
+  }
+  if (chroma_stride == 0) {
+    chroma_stride = luma_stride;
+  }
+  if (chroma_data == nullptr) {
+    chroma_data = &reinterpret_cast<uint16_t*>(image->data)[image->luma_stride * image->height];
+  }
 
-  uint16_t y_uint = reinterpret_cast<uint16_t*>(image->data)[pixel_y_idx]
-                  >> 6;
-  uint16_t u_uint = reinterpret_cast<uint16_t*>(image->data)[pixel_count + pixel_uv_idx * 2]
-                  >> 6;
-  uint16_t v_uint = reinterpret_cast<uint16_t*>(image->data)[pixel_count + pixel_uv_idx * 2 + 1]
-                  >> 6;
+  size_t pixel_y_idx = y * luma_stride + x;
+  size_t pixel_u_idx = (y >> 1) * chroma_stride + (x & ~0x1);
+  size_t pixel_v_idx = pixel_u_idx + 1;
+
+  uint16_t y_uint = luma_data[pixel_y_idx] >> 6;
+  uint16_t u_uint = chroma_data[pixel_u_idx] >> 6;
+  uint16_t v_uint = chroma_data[pixel_v_idx] >> 6;
 
   // Conversions include taking narrow-range into account.
   return {{{ (static_cast<float>(y_uint) - 64.0f) / 876.0f,
diff --git a/libs/jpegrecoverymap/tests/data/raw_p010_image_with_stride.p010 b/libs/jpegrecoverymap/tests/data/raw_p010_image_with_stride.p010
new file mode 100644
index 0000000..e7a5dc8
--- /dev/null
+++ b/libs/jpegrecoverymap/tests/data/raw_p010_image_with_stride.p010
Binary files differ
diff --git a/libs/jpegrecoverymap/tests/jpegr_test.cpp b/libs/jpegrecoverymap/tests/jpegr_test.cpp
index 7c669ab..229d7dc 100644
--- a/libs/jpegrecoverymap/tests/jpegr_test.cpp
+++ b/libs/jpegrecoverymap/tests/jpegr_test.cpp
@@ -24,10 +24,12 @@
 #include <utils/Log.h>
 
 #define RAW_P010_IMAGE "/sdcard/Documents/raw_p010_image.p010"
+#define RAW_P010_IMAGE_WITH_STRIDE "/sdcard/Documents/raw_p010_image_with_stride.p010"
 #define RAW_YUV420_IMAGE "/sdcard/Documents/raw_yuv420_image.yuv420"
 #define JPEG_IMAGE "/sdcard/Documents/jpeg_image.jpg"
 #define TEST_IMAGE_WIDTH 1280
 #define TEST_IMAGE_HEIGHT 720
+#define TEST_IMAGE_STRIDE 1288
 #define DEFAULT_JPEG_QUALITY 90
 
 #define SAVE_ENCODING_RESULT true
@@ -97,6 +99,7 @@
   virtual void TearDown();
 
   struct jpegr_uncompressed_struct mRawP010Image;
+  struct jpegr_uncompressed_struct mRawP010ImageWithStride;
   struct jpegr_uncompressed_struct mRawYuv420Image;
   struct jpegr_compressed_struct mJpegImage;
 };
@@ -107,6 +110,7 @@
 void JpegRTest::SetUp() {}
 void JpegRTest::TearDown() {
   free(mRawP010Image.data);
+  free(mRawP010ImageWithStride.data);
   free(mRawYuv420Image.data);
   free(mJpegImage.data);
 }
@@ -249,6 +253,61 @@
   free(decodedJpegR.data);
 }
 
+/* Test Encode API-0 (with stride) and decode */
+TEST_F(JpegRTest, encodeFromP010WithStrideThenDecode) {
+  int ret;
+
+  // Load input files.
+  if (!loadFile(RAW_P010_IMAGE_WITH_STRIDE, mRawP010ImageWithStride.data, nullptr)) {
+    FAIL() << "Load file " << RAW_P010_IMAGE_WITH_STRIDE << " failed";
+  }
+  mRawP010ImageWithStride.width = TEST_IMAGE_WIDTH;
+  mRawP010ImageWithStride.height = TEST_IMAGE_HEIGHT;
+  mRawP010ImageWithStride.luma_stride = TEST_IMAGE_STRIDE;
+  mRawP010ImageWithStride.colorGamut = jpegr_color_gamut::JPEGR_COLORGAMUT_BT2100;
+
+  JpegR jpegRCodec;
+
+  jpegr_compressed_struct jpegR;
+  jpegR.maxLength = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * sizeof(uint8_t);
+  jpegR.data = malloc(jpegR.maxLength);
+  ret = jpegRCodec.encodeJPEGR(
+      &mRawP010ImageWithStride, jpegr_transfer_function::JPEGR_TF_HLG, &jpegR,
+      DEFAULT_JPEG_QUALITY, nullptr);
+  if (ret != OK) {
+    FAIL() << "Error code is " << ret;
+  }
+  if (SAVE_ENCODING_RESULT) {
+    // Output image data to file
+    std::string filePath = "/sdcard/Documents/encoded_from_p010_input.jpgr";
+    std::ofstream imageFile(filePath.c_str(), std::ofstream::binary);
+    if (!imageFile.is_open()) {
+      ALOGE("%s: Unable to create file %s", __FUNCTION__, filePath.c_str());
+    }
+    imageFile.write((const char*)jpegR.data, jpegR.length);
+  }
+
+  jpegr_uncompressed_struct decodedJpegR;
+  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 8;
+  decodedJpegR.data = malloc(decodedJpegRSize);
+  ret = jpegRCodec.decodeJPEGR(&jpegR, &decodedJpegR);
+  if (ret != OK) {
+    FAIL() << "Error code is " << ret;
+  }
+  if (SAVE_DECODING_RESULT) {
+    // Output image data to file
+    std::string filePath = "/sdcard/Documents/decoded_from_p010_input.rgb";
+    std::ofstream imageFile(filePath.c_str(), std::ofstream::binary);
+    if (!imageFile.is_open()) {
+      ALOGE("%s: Unable to create file %s", __FUNCTION__, filePath.c_str());
+    }
+    imageFile.write((const char*)decodedJpegR.data, decodedJpegRSize);
+  }
+
+  free(jpegR.data);
+  free(decodedJpegR.data);
+}
+
 /* Test Encode API-1 and decode */
 TEST_F(JpegRTest, encodeFromRawHdrAndSdrThenDecode) {
   int ret;