jpegr: make output color format configurable for decoder

Test: jpegr_test
Bug: b/264715926
Change-Id: Id3fc2296cb0a82628f52724461ec4a19d898a48f
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegr.h b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegr.h
index 5455ba6..a433e8a 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegr.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegr.h
@@ -38,6 +38,14 @@
   JPEGR_TF_SRGB = 3,
 } jpegr_transfer_function;
 
+// Target output formats for decoder
+typedef enum {
+  JPEGR_OUTPUT_SDR,          // SDR in RGBA_8888 color format
+  JPEGR_OUTPUT_HDR_LINEAR,   // HDR in F16 color format (linear)
+  JPEGR_OUTPUT_HDR_PQ,       // HDR in RGBA_1010102 color format (PQ transfer function)
+  JPEGR_OUTPUT_HDR_HLG,      // HDR in RGBA_1010102 color format (HLG transfer function)
+} jpegr_output_format;
+
 struct jpegr_info_struct {
     size_t width;
     size_t height;
@@ -195,20 +203,15 @@
      * @param compressed_jpegr_image compressed JPEGR image
      * @param dest destination of the uncompressed JPEGR image
      * @param exif destination of the decoded EXIF metadata.
-     * @param request_sdr flag that request SDR output. If set to true, decoder will only decode
-     *                    the primary image which is SDR. Setting of request_sdr and input source
-     *                    (HDR or SDR) can be found in the table below:
-     *                    |  input source  |  request_sdr  |  output of decoding  |
-     *                    |       HDR      |     true      |          SDR         |
-     *                    |       HDR      |     false     |          HDR         |
-     *                    |       SDR      |     true      |          SDR         |
-     *                    |       SDR      |     false     |          SDR         |
+     * @param output_format flag for setting output color format. if set to
+     *                      {@code JPEGR_OUTPUT_SDR}, decoder will only decode the primary image
+     *                      which is SDR. Default value is JPEGR_OUTPUT_HDR_LINEAR.
      * @return NO_ERROR if decoding succeeds, error code if error occurs.
      */
     status_t decodeJPEGR(jr_compressed_ptr compressed_jpegr_image,
                          jr_uncompressed_ptr dest,
                          jr_exif_ptr exif = nullptr,
-                         bool request_sdr = false);
+                         jpegr_output_format output_format = JPEGR_OUTPUT_HDR_LINEAR);
 
     /*
     * Gets Info from JPEGR file without decoding it.
@@ -249,12 +252,16 @@
      * @param uncompressed_yuv_420_image uncompressed SDR image in YUV_420 color format
      * @param uncompressed_recovery_map uncompressed recovery map
      * @param metadata JPEG/R metadata extracted from XMP.
+     * @param output_format flag for setting output color format. if set to
+     *                      {@code JPEGR_OUTPUT_SDR}, decoder will only decode the primary image
+     *                      which is SDR. Default value is JPEGR_OUTPUT_HDR_LINEAR.
      * @param dest reconstructed HDR image
      * @return NO_ERROR if calculation succeeds, error code if error occurs.
      */
     status_t applyRecoveryMap(jr_uncompressed_ptr uncompressed_yuv_420_image,
                               jr_uncompressed_ptr uncompressed_recovery_map,
                               jr_metadata_ptr metadata,
+                              jpegr_output_format output_format,
                               jr_uncompressed_ptr dest);
 
 private:
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymapmath.h b/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymapmath.h
index c12cee9..8b5318f 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymapmath.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymapmath.h
@@ -115,6 +115,14 @@
   return temp /= rhs;
 }
 
+inline uint16_t floatToHalf(float f) {
+  uint32_t x = *((uint32_t*)&f);
+  uint16_t h = ((x >> 16) & 0x8000)
+             | ((((x & 0x7f800000) - 0x38000000) >> 13) & 0x7c00)
+             | ((x >> 13) & 0x03ff);
+  return h;
+}
+
 constexpr size_t kRecoveryFactorPrecision = 10;
 constexpr size_t kRecoveryFactorNumEntries = 1 << kRecoveryFactorPrecision;
 struct RecoveryLUT {
@@ -392,6 +400,13 @@
  */
 uint32_t colorToRgba1010102(Color e_gamma);
 
+/*
+ * Convert from Color to F16.
+ *
+ * Alpha always set to 1.0.
+ */
+uint64_t colorToRgbaF16(Color e_gamma);
+
 } // namespace android::jpegrecoverymap
 
 #endif // ANDROID_JPEGRECOVERYMAP_RECOVERYMAPMATH_H
diff --git a/libs/jpegrecoverymap/jpegr.cpp b/libs/jpegrecoverymap/jpegr.cpp
index c22020a..79b1ae3 100644
--- a/libs/jpegrecoverymap/jpegr.cpp
+++ b/libs/jpegrecoverymap/jpegr.cpp
@@ -351,14 +351,14 @@
 status_t JpegR::decodeJPEGR(jr_compressed_ptr compressed_jpegr_image,
                             jr_uncompressed_ptr dest,
                             jr_exif_ptr exif,
-                            bool request_sdr) {
+                            jpegr_output_format output_format) {
   if (compressed_jpegr_image == nullptr || dest == nullptr) {
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
   // TODO: fill EXIF data
   (void) exif;
 
-  if (request_sdr) {
+  if (output_format == JPEGR_OUTPUT_SDR) {
     JpegDecoderHelper jpeg_decoder;
     if (!jpeg_decoder.decompressImage(compressed_jpegr_image->data, compressed_jpegr_image->length,
                                       true)) {
@@ -404,7 +404,7 @@
     return ERROR_JPEGR_DECODE_ERROR;
   }
 
-  JPEGR_CHECK(applyRecoveryMap(&uncompressed_yuv_420_image, &map, &metadata, dest));
+  JPEGR_CHECK(applyRecoveryMap(&uncompressed_yuv_420_image, &map, &metadata, output_format, dest));
   return NO_ERROR;
 }
 
@@ -639,6 +639,7 @@
 status_t JpegR::applyRecoveryMap(jr_uncompressed_ptr uncompressed_yuv_420_image,
                                  jr_uncompressed_ptr uncompressed_recovery_map,
                                  jr_metadata_ptr metadata,
+                                 jpegr_output_format output_format,
                                  jr_uncompressed_ptr dest) {
   if (uncompressed_yuv_420_image == nullptr
    || uncompressed_recovery_map == nullptr
@@ -654,18 +655,12 @@
 
   JobQueue jobQueue;
   std::function<void()> applyRecMap = [uncompressed_yuv_420_image, uncompressed_recovery_map,
-                                       metadata, dest, &jobQueue, &idwTable,
+                                       metadata, dest, &jobQueue, &idwTable, output_format,
                                        &recoveryLUT]() -> void {
     const float hdr_ratio = metadata->maxContentBoost;
     size_t width = uncompressed_yuv_420_image->width;
     size_t height = uncompressed_yuv_420_image->height;
 
-#if USE_HLG_OETF_LUT
-    ColorTransformFn hdrOetf = hlgOetfLUT;
-#else
-    ColorTransformFn hdrOetf = hlgOetf;
-#endif
-
     size_t rowStart, rowEnd;
     while (jobQueue.dequeueJob(rowStart, rowEnd)) {
       for (size_t y = rowStart; y < rowEnd; ++y) {
@@ -693,11 +688,44 @@
 #else
           Color rgb_hdr = applyRecovery(rgb_sdr, recovery, metadata);
 #endif
-          Color rgb_gamma_hdr = hdrOetf(rgb_hdr / metadata->maxContentBoost);
-          uint32_t rgba1010102 = colorToRgba1010102(rgb_gamma_hdr);
-
+          rgb_hdr = rgb_hdr / metadata->maxContentBoost;
           size_t pixel_idx = x + y * width;
-          reinterpret_cast<uint32_t*>(dest->data)[pixel_idx] = rgba1010102;
+
+          switch (output_format) {
+            case JPEGR_OUTPUT_HDR_LINEAR:
+            {
+              uint64_t rgba_f16 = colorToRgbaF16(rgb_hdr);
+              reinterpret_cast<uint64_t*>(dest->data)[pixel_idx] = rgba_f16;
+              break;
+            }
+            case JPEGR_OUTPUT_HDR_HLG:
+            {
+#if USE_HLG_OETF_LUT
+              ColorTransformFn hdrOetf = hlgOetfLUT;
+#else
+              ColorTransformFn hdrOetf = hlgOetf;
+#endif
+              Color rgb_gamma_hdr = hdrOetf(rgb_hdr);
+              uint32_t rgba_1010102 = colorToRgba1010102(rgb_gamma_hdr);
+              reinterpret_cast<uint32_t*>(dest->data)[pixel_idx] = rgba_1010102;
+              break;
+            }
+            case JPEGR_OUTPUT_HDR_PQ:
+            {
+#if USE_HLG_OETF_LUT
+              ColorTransformFn hdrOetf = pqOetfLUT;
+#else
+              ColorTransformFn hdrOetf = pqOetf;
+#endif
+              Color rgb_gamma_hdr = hdrOetf(rgb_hdr);
+              uint32_t rgba_1010102 = colorToRgba1010102(rgb_gamma_hdr);
+              reinterpret_cast<uint32_t*>(dest->data)[pixel_idx] = rgba_1010102;
+              break;
+            }
+            default:
+            {}
+              // Should be impossible to hit after input validation.
+          }
         }
       }
     }
diff --git a/libs/jpegrecoverymap/recoverymapmath.cpp b/libs/jpegrecoverymap/recoverymapmath.cpp
index 7812e18..20c32ed 100644
--- a/libs/jpegrecoverymap/recoverymapmath.cpp
+++ b/libs/jpegrecoverymap/recoverymapmath.cpp
@@ -631,4 +631,11 @@
        | (0x3 << 30);  // Set alpha to 1.0
 }
 
+uint64_t colorToRgbaF16(Color e_gamma) {
+  return (uint64_t) floatToHalf(e_gamma.r)
+       | (((uint64_t) floatToHalf(e_gamma.g)) << 16)
+       | (((uint64_t) floatToHalf(e_gamma.b)) << 32)
+       | (((uint64_t) floatToHalf(1.0f)) << 48);
+}
+
 } // namespace android::jpegrecoverymap
diff --git a/libs/jpegrecoverymap/tests/jpegr_test.cpp b/libs/jpegrecoverymap/tests/jpegr_test.cpp
index df212e1..0a7d20a 100644
--- a/libs/jpegrecoverymap/tests/jpegr_test.cpp
+++ b/libs/jpegrecoverymap/tests/jpegr_test.cpp
@@ -152,7 +152,7 @@
 
   timerStart(&applyRecMapTime);
   for (auto i = 0; i < kProfileCount; i++) {
-      ASSERT_EQ(OK, applyRecoveryMap(yuv420Image, map, metadata, dest));
+      ASSERT_EQ(OK, applyRecoveryMap(yuv420Image, map, metadata, JPEGR_OUTPUT_HDR_HLG, dest));
   }
   timerStop(&applyRecMapTime);
 
@@ -170,7 +170,7 @@
   jpegRCodec.encodeJPEGR(nullptr, nullptr, nullptr, static_cast<jpegr_transfer_function>(0),
                          nullptr);
   jpegRCodec.encodeJPEGR(nullptr, nullptr, static_cast<jpegr_transfer_function>(0), nullptr);
-  jpegRCodec.decodeJPEGR(nullptr, nullptr, nullptr, false);
+  jpegRCodec.decodeJPEGR(nullptr, nullptr, nullptr);
 }
 
 TEST_F(JpegRTest, writeXmpThenRead) {
@@ -228,7 +228,7 @@
   }
 
   jpegr_uncompressed_struct decodedJpegR;
-  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 4;
+  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 8;
   decodedJpegR.data = malloc(decodedJpegRSize);
   ret = jpegRCodec.decodeJPEGR(&jpegR, &decodedJpegR);
   if (ret != OK) {
@@ -236,7 +236,7 @@
   }
   if (SAVE_DECODING_RESULT) {
     // Output image data to file
-    std::string filePath = "/sdcard/Documents/decoded_from_p010_input.rgb10";
+    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());
@@ -289,7 +289,7 @@
   }
 
   jpegr_uncompressed_struct decodedJpegR;
-  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 4;
+  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 8;
   decodedJpegR.data = malloc(decodedJpegRSize);
   ret = jpegRCodec.decodeJPEGR(&jpegR, &decodedJpegR);
   if (ret != OK) {
@@ -297,7 +297,7 @@
   }
   if (SAVE_DECODING_RESULT) {
     // Output image data to file
-    std::string filePath = "/sdcard/Documents/decoded_from_p010_yuv420p_input.rgb10";
+    std::string filePath = "/sdcard/Documents/decoded_from_p010_yuv420p_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());
@@ -354,7 +354,7 @@
   }
 
   jpegr_uncompressed_struct decodedJpegR;
-  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 4;
+  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 8;
   decodedJpegR.data = malloc(decodedJpegRSize);
   ret = jpegRCodec.decodeJPEGR(&jpegR, &decodedJpegR);
   if (ret != OK) {
@@ -362,7 +362,7 @@
   }
   if (SAVE_DECODING_RESULT) {
     // Output image data to file
-    std::string filePath = "/sdcard/Documents/decoded_from_p010_yuv420p_jpeg_input.rgb10";
+    std::string filePath = "/sdcard/Documents/decoded_from_p010_yuv420p_jpeg_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());
@@ -435,7 +435,7 @@
   }
 
   jpegr_uncompressed_struct decodedJpegR;
-  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 4;
+  int decodedJpegRSize = TEST_IMAGE_WIDTH * TEST_IMAGE_HEIGHT * 8;
   decodedJpegR.data = malloc(decodedJpegRSize);
   ret = jpegRCodec.decodeJPEGR(&jpegR, &decodedJpegR);
   if (ret != OK) {
@@ -443,7 +443,7 @@
   }
   if (SAVE_DECODING_RESULT) {
     // Output image data to file
-    std::string filePath = "/sdcard/Documents/decoded_from_p010_jpeg_input.rgb10";
+    std::string filePath = "/sdcard/Documents/decoded_from_p010_jpeg_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());