jpegrecoverymap: implement tonemap() method

Bug: b/252835416
Test: recoverymap_test.cpp#encodeFromP010ThenDecode
Change-Id: I57738de8a9c78a3f9dcf349fe5bb3604d3a5eb27
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymap.h b/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymap.h
index 905bf16..05b1421 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymap.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymap.h
@@ -339,11 +339,11 @@
     /*
      * This method will tone map a HDR image to an SDR image.
      *
-     * @param uncompressed_p010_image (input) uncompressed P010 image
+     * @param src (input) uncompressed P010 image
      * @param dest (output) tone mapping result as a YUV_420 image
      * @return NO_ERROR if calculation succeeds, error code if error occurs.
      */
-    status_t toneMap(jr_uncompressed_ptr uncompressed_p010_image,
+    status_t toneMap(jr_uncompressed_ptr src,
                      jr_uncompressed_ptr dest);
 };
 
diff --git a/libs/jpegrecoverymap/recoverymap.cpp b/libs/jpegrecoverymap/recoverymap.cpp
index 53fa8ce..f362e00 100644
--- a/libs/jpegrecoverymap/recoverymap.cpp
+++ b/libs/jpegrecoverymap/recoverymap.cpp
@@ -25,7 +25,6 @@
 #include <image_io/jpeg/jpeg_scanner.h>
 #include <image_io/jpeg/jpeg_info_builder.h>
 #include <image_io/base/data_segment_data_source.h>
-#include <utils/Log.h>
 
 #include <memory>
 #include <sstream>
@@ -214,6 +213,9 @@
   }
 
   jpegr_uncompressed_struct uncompressed_yuv_420_image;
+  unique_ptr<uint8_t[]> uncompressed_yuv_420_image_data = make_unique<uint8_t[]>(
+      uncompressed_p010_image->width * uncompressed_p010_image->height * 3 / 2);
+  uncompressed_yuv_420_image.data = uncompressed_yuv_420_image_data.get();
   JPEGR_CHECK(toneMap(uncompressed_p010_image, &uncompressed_yuv_420_image));
 
   jpegr_uncompressed_struct map;
@@ -240,7 +242,7 @@
   jpeg.length = jpeg_encoder.getCompressedImageSize();
 
   jpegr_exif_struct new_exif;
-  if (exif->data == nullptr) {
+  if (exif == nullptr || exif->data == nullptr) {
       new_exif.length = PSEUDO_EXIF_PACKAGE_LENGTH;
   } else {
       new_exif.length = exif->length + EXIF_J_R_ENTRY_LENGTH;
@@ -889,18 +891,39 @@
   return NO_ERROR;
 }
 
-status_t RecoveryMap::toneMap(jr_uncompressed_ptr uncompressed_p010_image,
+status_t RecoveryMap::toneMap(jr_uncompressed_ptr src,
                               jr_uncompressed_ptr dest) {
-  if (uncompressed_p010_image == nullptr || dest == nullptr) {
+  if (src == nullptr || dest == nullptr) {
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  dest->width = uncompressed_p010_image->width;
-  dest->height = uncompressed_p010_image->height;
-  unique_ptr<uint8_t[]> dest_data = make_unique<uint8_t[]>(dest->width * dest->height * 3 / 2);
-  dest->data = dest_data.get();
+  dest->width = src->width;
+  dest->height = src->height;
 
-  // TODO: Tone map algorighm here.
+  size_t pixel_count = src->width * src->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);
+
+      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;
+
+      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];
+
+      *y = static_cast<uint8_t>((y_uint >> 2) & 0xff);
+      *u = static_cast<uint8_t>((u_uint >> 2) & 0xff);
+      *v = static_cast<uint8_t>((v_uint >> 2) & 0xff);
+    }
+  }
+
+  dest->colorGamut = src->colorGamut;
 
   return NO_ERROR;
 }
diff --git a/libs/jpegrecoverymap/tests/recoverymap_test.cpp b/libs/jpegrecoverymap/tests/recoverymap_test.cpp
index c3c6fd4..8ff12fb 100644
--- a/libs/jpegrecoverymap/tests/recoverymap_test.cpp
+++ b/libs/jpegrecoverymap/tests/recoverymap_test.cpp
@@ -114,58 +114,57 @@
 }
 
 /* Test Encode API-0 and decode */
-// TODO: enable when tonemapper is ready.
-//TEST_F(RecoveryMapTest, encodeFromP010ThenDecode) {
-//  int ret;
-//
-//  // Load input files.
-//  if (!loadFile(RAW_P010_IMAGE, mRawP010Image.data, nullptr)) {
-//    FAIL() << "Load file " << RAW_P010_IMAGE << " failed";
-//  }
-//  mRawP010Image.width = TEST_IMAGE_WIDTH;
-//  mRawP010Image.height = TEST_IMAGE_HEIGHT;
-//  mRawP010Image.colorGamut = jpegr_color_gamut::JPEGR_COLORGAMUT_BT2100;
-//
-//  RecoveryMap recoveryMap;
-//
-//  jpegr_compressed_struct jpegR;
-//  jpegR.maxLength = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * sizeof(uint8_t);
-//  jpegR.data = malloc(jpegR.maxLength);
-//  ret = recoveryMap.encodeJPEGR(
-//      &mRawP010Image, jpegr_transfer_function::JPEGR_TF_HLG, &jpegR, 90, 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_jpeg_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 * 4;
-//  decodedJpegR.data = malloc(decodedJpegRSize);
-//  ret = recoveryMap.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_jpeg_input.rgb10";
-//    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_F(RecoveryMapTest, encodeFromP010ThenDecode) {
+  int ret;
+
+  // Load input files.
+  if (!loadFile(RAW_P010_IMAGE, mRawP010Image.data, nullptr)) {
+    FAIL() << "Load file " << RAW_P010_IMAGE << " failed";
+  }
+  mRawP010Image.width = TEST_IMAGE_WIDTH;
+  mRawP010Image.height = TEST_IMAGE_HEIGHT;
+  mRawP010Image.colorGamut = jpegr_color_gamut::JPEGR_COLORGAMUT_BT2100;
+
+  RecoveryMap recoveryMap;
+
+  jpegr_compressed_struct jpegR;
+  jpegR.maxLength = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * sizeof(uint8_t);
+  jpegR.data = malloc(jpegR.maxLength);
+  ret = recoveryMap.encodeJPEGR(
+      &mRawP010Image, 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_jpeg_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 * 4;
+  decodedJpegR.data = malloc(decodedJpegRSize);
+  ret = recoveryMap.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_jpeg_input.rgb10";
+    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(RecoveryMapTest, encodeFromRawHdrAndSdrThenDecode) {