jpegrecoverymap: add initial recovery map calculations.

This change adds the starting point for generating and applying the
recovery map. A follow-up change will add more robust tests for this
update (eg. unit testing color conversions).

There are a few other known TODOs remaining for these map operations:
  * Clean up handling around hdr_ratio (ie. utilizing XMP)
  * Add color space information for inputs and utilize it for ICC data
  * Add handling for PQ encode/decode

No-Typo-Check: Lint suggesting typo in code as if it's a comment
Test: build
Bug: b/252835416
Change-Id: I2f6a1bf3b046036292afe46bbd2396a87cdc5164

11 files changed, 584 insertions, 51 deletions

diff --git a/libs/jpegrecoverymap/Android.bp b/libs/jpegrecoverymap/Android.bp
index a9248ab2e8..0375915839 100644
--- a/libs/jpegrecoverymap/Android.bp
+++ b/libs/jpegrecoverymap/Android.bp
@@ -30,10 +30,12 @@ cc_library_static {
 
     srcs: [
         "recoverymap.cpp",
+        "recoverymapmath.cpp",
     ],
 
     shared_libs: [
         "libimage_io",
+        "libjpeg",
         "libutils",
     ],
 }
@@ -64,4 +66,4 @@ cc_library_static {
     srcs: [
         "jpegdecoder.cpp",
     ],
-}
\ No newline at end of file
+}
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegdecoder.h b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegdecoder.h
index 2ab75503a5..df24b10ebc 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegdecoder.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegdecoder.h
@@ -14,6 +14,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+
+#ifndef ANDROID_JPEGRECOVERYMAP_JPEGDECODER_H
+#define ANDROID_JPEGRECOVERYMAP_JPEGDECODER_H
+
 // We must include cstdio before jpeglib.h. It is a requirement of libjpeg.
 #include <cstdio>
 extern "C" {
@@ -41,12 +45,22 @@ public:
      * Returns the decompressed raw image buffer pointer. This method must be called only after
      * calling decompressImage().
      */
-    const void* getDecompressedImagePtr();
+    void* getDecompressedImagePtr();
     /*
      * Returns the decompressed raw image buffer size. This method must be called only after
      * calling decompressImage().
      */
     size_t getDecompressedImageSize();
+    /*
+     * Returns the image width in pixels. This method must be called only after calling
+     * decompressImage().
+     */
+    size_t getDecompressedImageWidth();
+    /*
+     * Returns the image width in pixels. This method must be called only after calling
+     * decompressImage().
+     */
+    size_t getDecompressedImageHeight();
 private:
     bool decode(const void* image, int length);
     // Returns false if errors occur.
@@ -56,7 +70,12 @@ private:
     // Process 16 lines of Y and 16 lines of U/V each time.
     // We must pass at least 16 scanlines according to libjpeg documentation.
     static const int kCompressBatchSize = 16;
-    // The buffer that holds the compressed result.
+    // The buffer that holds the decompressed result.
     std::vector<JOCTET> mResultBuffer;
+    // Resolution of the decompressed image.
+    size_t mWidth;
+    size_t mHeight;
 };
 } /* namespace android  */
+
+#endif // ANDROID_JPEGRECOVERYMAP_JPEGDECODER_H
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegencoder.h b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegencoder.h
index 9641fda24c..61aeb8ace7 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegencoder.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegencoder.h
@@ -14,6 +14,9 @@
  * limitations under the License.
  */
 
+#ifndef ANDROID_JPEGRECOVERYMAP_JPEGENCODER_H
+#define ANDROID_JPEGRECOVERYMAP_JPEGENCODER_H
+
 // We must include cstdio before jpeglib.h. It is a requirement of libjpeg.
 #include <cstdio>
 
@@ -50,7 +53,7 @@ public:
      * Returns the compressed JPEG buffer pointer. This method must be called only after calling
      * compressImage().
      */
-    const void* getCompressedImagePtr();
+    void* getCompressedImagePtr();
 
     /*
      * Returns the compressed JPEG buffer size. This method must be called only after calling
@@ -87,4 +90,6 @@ private:
     std::vector<JOCTET> mResultBuffer;
 };
 
-} /* namespace android  */
\ No newline at end of file
+} /* namespace android  */
+
+#endif // ANDROID_JPEGRECOVERYMAP_JPEGENCODER_H
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegrerrorcode.h b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegrerrorcode.h
index 49ab34d154..194cd2f403 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/jpegrerrorcode.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/jpegrerrorcode.h
@@ -35,6 +35,8 @@ enum {
     ERROR_JPEGR_INVALID_INPUT_TYPE      = JPEGR_IO_ERROR_BASE,
     ERROR_JPEGR_INVALID_OUTPUT_TYPE     = JPEGR_IO_ERROR_BASE - 1,
     ERROR_JPEGR_INVALID_NULL_PTR        = JPEGR_IO_ERROR_BASE - 2,
+    ERROR_JPEGR_RESOLUTION_MISMATCH     = JPEGR_IO_ERROR_BASE - 3,
+    ERROR_JPEGR_BUFFER_TOO_SMALL        = JPEGR_IO_ERROR_BASE - 4,
 
     JPEGR_RUNTIME_ERROR_BASE            = -20000,
     ERROR_JPEGR_ENCODE_ERROR            = JPEGR_RUNTIME_ERROR_BASE - 1,
@@ -43,4 +45,4 @@ enum {
     ERROR_JPEGR_METADATA_ERROR          = JPEGR_RUNTIME_ERROR_BASE - 4,
 };
 
-}  // namespace android::recoverymap
\ No newline at end of file
+}  // namespace android::recoverymap
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymap.h b/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymap.h
index 06b2ab5874..94b7543e14 100644
--- a/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymap.h
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymap.h
@@ -14,6 +14,9 @@
  * limitations under the License.
  */
 
+#ifndef ANDROID_JPEGRECOVERYMAP_RECOVERYMAP_H
+#define ANDROID_JPEGRECOVERYMAP_RECOVERYMAP_H
+
 #include "jpegrerrorcode.h"
 
 namespace android::recoverymap {
@@ -71,7 +74,8 @@ public:
      * Compress JPEGR image from 10-bit HDR YUV and 8-bit SDR YUV.
      *
      * Generate recovery map from the HDR and SDR inputs, compress SDR YUV to 8-bit JPEG and append
-     * the recovery map to the end of the compressed JPEG.
+     * the recovery map to the end of the compressed JPEG. HDR and SDR inputs must be the same
+     * resolution and color space.
      * @param uncompressed_p010_image uncompressed HDR image in P010 color format
      * @param uncompressed_yuv_420_image uncompressed SDR image in YUV_420 color format
      * @param dest destination of the compressed JPEGR image
@@ -84,7 +88,7 @@ public:
      */
     status_t encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
                          jr_uncompressed_ptr uncompressed_yuv_420_image,
-                         void* dest,
+                         jr_compressed_ptr dest,
                          int quality,
                          jr_exif_ptr exif,
                          float hdr_ratio = 0.0f);
@@ -95,7 +99,7 @@ public:
      * This method requires HAL Hardware JPEG encoder.
      *
      * Generate recovery map from the HDR and SDR inputs, append the recovery map to the end of the
-     * compressed JPEG.
+     * compressed JPEG. HDR and SDR inputs must be the same resolution and color space.
      * @param uncompressed_p010_image uncompressed HDR image in P010 color format
      * @param uncompressed_yuv_420_image uncompressed SDR image in YUV_420 color format
      * @param compressed_jpeg_image compressed 8-bit JPEG image
@@ -106,8 +110,8 @@ public:
      */
     status_t encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
                          jr_uncompressed_ptr uncompressed_yuv_420_image,
-                         void* compressed_jpeg_image,
-                         void* dest,
+                         jr_compressed_ptr compressed_jpeg_image,
+                         jr_compressed_ptr dest,
                          float hdr_ratio = 0.0f);
 
     /*
@@ -116,7 +120,8 @@ public:
      * This method requires HAL Hardware JPEG encoder.
      *
      * Decode the compressed 8-bit JPEG image to YUV SDR, generate recovery map from the HDR input
-     * and the decoded SDR result, append the recovery map to the end of the compressed JPEG.
+     * and the decoded SDR result, append the recovery map to the end of the compressed JPEG. HDR
+     * and SDR inputs must be the same resolution and color space.
      * @param uncompressed_p010_image uncompressed HDR image in P010 color format
      * @param compressed_jpeg_image compressed 8-bit JPEG image
      * @param dest destination of the compressed JPEGR image
@@ -125,8 +130,8 @@ public:
      * @return NO_ERROR if encoding succeeds, error code if error occurs.
      */
     status_t encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
-                         void* compressed_jpeg_image,
-                         void* dest,
+                         jr_compressed_ptr compressed_jpeg_image,
+                         jr_compressed_ptr dest,
                          float hdr_ratio = 0.0f);
 
     /*
@@ -147,7 +152,7 @@ public:
      *                    |       SDR      |     false     |          SDR         |
      * @return NO_ERROR if decoding succeeds, error code if error occurs.
      */
-    status_t decodeJPEGR(void* compressed_jpegr_image,
+    status_t decodeJPEGR(jr_compressed_ptr compressed_jpegr_image,
                          jr_uncompressed_ptr dest,
                          jr_exif_ptr exif = nullptr,
                          bool request_sdr = false);
@@ -159,7 +164,7 @@ private:
      * @param dest decoded recover map
      * @return NO_ERROR if decoding succeeds, error code if error occurs.
      */
-    status_t decodeRecoveryMap(jr_compressed_ptr compressed_recovery_map,
+    status_t decompressRecoveryMap(jr_compressed_ptr compressed_recovery_map,
                                jr_uncompressed_ptr dest);
 
     /*
@@ -169,16 +174,17 @@ private:
      * @param dest encoded recover map
      * @return NO_ERROR if encoding succeeds, error code if error occurs.
      */
-    status_t encodeRecoveryMap(jr_uncompressed_ptr uncompressed_recovery_map,
+    status_t compressRecoveryMap(jr_uncompressed_ptr uncompressed_recovery_map,
                                jr_compressed_ptr dest);
 
     /*
      * This method is called in the encoding pipeline. It will take the uncompressed 8-bit and
-     * 10-bit yuv images as input, and calculate the uncompressed recovery map.
+     * 10-bit yuv images as input, and calculate the uncompressed recovery map. The input images
+     * must be the same resolution.
      *
      * @param uncompressed_yuv_420_image uncompressed SDR image in YUV_420 color format
      * @param uncompressed_p010_image uncompressed HDR image in P010 color format
-     * @param dest recover map
+     * @param dest recovery map; caller responsible for memory of data
      * @return NO_ERROR if calculation succeeds, error code if error occurs.
      */
     status_t generateRecoveryMap(jr_uncompressed_ptr uncompressed_yuv_420_image,
@@ -207,7 +213,7 @@ private:
      * @param dest destination of compressed recovery map
      * @return NO_ERROR if calculation succeeds, error code if error occurs.
      */
-    status_t extractRecoveryMap(void* compressed_jpegr_image, jr_compressed_ptr dest);
+    status_t extractRecoveryMap(jr_compressed_ptr compressed_jpegr_image, jr_compressed_ptr dest);
 
     /*
      * This method is called in the encoding pipeline. It will take the standard 8-bit JPEG image
@@ -219,9 +225,9 @@ private:
      * @param dest compressed JPEGR image
      * @return NO_ERROR if calculation succeeds, error code if error occurs.
      */
-    status_t appendRecoveryMap(void* compressed_jpeg_image,
+    status_t appendRecoveryMap(jr_compressed_ptr compressed_jpeg_image,
                                jr_compressed_ptr compressed_recovery_map,
-                               void* dest);
+                               jr_compressed_ptr dest);
 
     /*
      * This method generates XMP metadata.
@@ -266,3 +272,5 @@ private:
 };
 
 } // namespace android::recoverymap
+
+#endif // ANDROID_JPEGRECOVERYMAP_RECOVERYMAP_H
diff --git a/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymapmath.h b/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymapmath.h
new file mode 100644
index 0000000000..8e9b07bf3d
--- /dev/null
+++ b/libs/jpegrecoverymap/include/jpegrecoverymap/recoverymapmath.h
@@ -0,0 +1,106 @@
+/*
+ * 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.
+ */
+
+#ifndef ANDROID_JPEGRECOVERYMAP_RECOVERYMAPMATH_H
+#define ANDROID_JPEGRECOVERYMAP_RECOVERYMAPMATH_H
+
+#include <stdint.h>
+
+#include <jpegrecoverymap/recoverymap.h>
+
+namespace android::recoverymap {
+
+const float kSdrWhiteNits = 100.0f;
+
+struct Color {
+  union {
+    struct {
+      float r;
+      float g;
+      float b;
+    };
+    struct {
+      float y;
+      float u;
+      float v;
+    };
+  };
+};
+
+/*
+ * Convert from OETF'd bt.2100 RGB to YUV, according to BT.2100
+ */
+Color bt2100RgbToYuv(Color e);
+
+/*
+ * Convert srgb YUV to RGB, according to ECMA TR/98.
+ */
+Color srgbYuvToRgb(Color e);
+
+/*
+ * TODO: better source for srgb transfer function
+ * Convert from srgb to linear, according to https://en.wikipedia.org/wiki/SRGB.
+ * [0.0, 1.0] range in and out.
+ */
+float srgbInvOetf(float e);
+Color srgbInvOetf(Color e);
+
+/*
+ * Convert from HLG to scene luminance in nits, according to BT.2100.
+ */
+float hlgInvOetf(float e);
+
+/*
+ * Convert from scene luminance in nits to HLG,  according to BT.2100.
+ */
+float hlgOetf(float e);
+Color hlgOetf(Color e);
+
+/*
+ * Calculate the 8-bit unsigned integer recovery value for the given SDR and HDR
+ * luminances in linear space, and the hdr ratio to encode against.
+ */
+uint8_t encodeRecovery(float y_sdr, float y_hdr, float hdr_ratio);
+
+/*
+ * Calculates the linear luminance in nits after applying the given recovery
+ * value, with the given hdr ratio, to the given sdr input in the range [0, 1].
+ */
+Color applyRecovery(Color e, float recovery, float hdr_ratio);
+
+/*
+ * Helper for sampling from images.
+ */
+Color getYuv420Pixel(jr_uncompressed_ptr image, size_t x, size_t y);
+
+/*
+ * Sample the recovery value for the map from a given x,y coordinate on a scale
+ * that is map scale factor larger than the map size.
+ */
+float sampleMap(jr_uncompressed_ptr map, size_t map_scale_factor, size_t x, size_t y);
+
+/*
+ * Sample the image Y value at the provided location, with a weighting based on nearby pixels
+ * and the map scale factor.
+ *
+ * Expect narrow-range image data for P010.
+ */
+float sampleYuv420Y(jr_uncompressed_ptr map, size_t map_scale_factor, size_t x, size_t y);
+float sampleP010Y(jr_uncompressed_ptr map, size_t map_scale_factor, size_t x, size_t y);
+
+} // namespace android::recoverymap
+
+#endif // ANDROID_JPEGRECOVERYMAP_RECOVERYMAPMATH_H
diff --git a/libs/jpegrecoverymap/jpegdecoder.cpp b/libs/jpegrecoverymap/jpegdecoder.cpp
index 22a5389648..c1fb6c3f1d 100644
--- a/libs/jpegrecoverymap/jpegdecoder.cpp
+++ b/libs/jpegrecoverymap/jpegdecoder.cpp
@@ -95,7 +95,7 @@ bool JpegDecoder::decompressImage(const void* image, int length) {
     return true;
 }
 
-const void* JpegDecoder::getDecompressedImagePtr() {
+void* JpegDecoder::getDecompressedImagePtr() {
     return mResultBuffer.data();
 }
 
@@ -103,6 +103,14 @@ size_t JpegDecoder::getDecompressedImageSize() {
     return mResultBuffer.size();
 }
 
+size_t JpegDecoder::getDecompressedImageWidth() {
+    return mWidth;
+}
+
+size_t JpegDecoder::getDecompressedImageHeight() {
+    return mHeight;
+}
+
 bool JpegDecoder::decode(const void* image, int length) {
     jpeg_decompress_struct cinfo;
     jpegr_source_mgr mgr(static_cast<const uint8_t*>(image), length);
@@ -119,6 +127,9 @@ bool JpegDecoder::decode(const void* image, int length) {
     cinfo.src = &mgr;
     jpeg_read_header(&cinfo, TRUE);
 
+    mWidth = cinfo.image_width;
+    mHeight = cinfo.image_height;
+
     if (cinfo.jpeg_color_space == JCS_YCbCr) {
         mResultBuffer.resize(cinfo.image_width * cinfo.image_height * 3 / 2, 0);
     } else if (cinfo.jpeg_color_space == JCS_GRAYSCALE) {
@@ -222,4 +233,4 @@ bool JpegDecoder::decompressSingleChannel(jpeg_decompress_struct* cinfo, const u
     return true;
 }
 
-} // namespace android
\ No newline at end of file
+} // namespace android
diff --git a/libs/jpegrecoverymap/jpegencoder.cpp b/libs/jpegrecoverymap/jpegencoder.cpp
index d45d9b33c9..1997bf9ea3 100644
--- a/libs/jpegrecoverymap/jpegencoder.cpp
+++ b/libs/jpegrecoverymap/jpegencoder.cpp
@@ -52,7 +52,7 @@ bool JpegEncoder::compressImage(const void* image, int width, int height, int qu
     return true;
 }
 
-const void* JpegEncoder::getCompressedImagePtr() {
+void* JpegEncoder::getCompressedImagePtr() {
     return mResultBuffer.data();
 }
 
@@ -236,4 +236,4 @@ bool JpegEncoder::compressSingleChannel(jpeg_compress_struct* cinfo, const uint8
     return true;
 }
 
-} // namespace android
\ No newline at end of file
+} // namespace android
diff --git a/libs/jpegrecoverymap/recoverymap.cpp b/libs/jpegrecoverymap/recoverymap.cpp
index 71e5f6f21a..4a90053b12 100644
--- a/libs/jpegrecoverymap/recoverymap.cpp
+++ b/libs/jpegrecoverymap/recoverymap.cpp
@@ -14,9 +14,20 @@
  * limitations under the License.
  */
 
-#include "image_io/xml/xml_writer.h"
+// TODO: need to clean up handling around hdr_ratio and passing it around
+// TODO: need to handle color space information; currently we assume everything
+// is srgb in.
+// TODO: handle PQ encode/decode (currently only HLG)
 
 #include <jpegrecoverymap/recoverymap.h>
+
+#include <jpegrecoverymap/jpegencoder.h>
+#include <jpegrecoverymap/jpegdecoder.h>
+#include <jpegrecoverymap/recoverymapmath.h>
+
+#include <image_io/xml/xml_writer.h>
+
+#include <memory>
 #include <sstream>
 #include <string>
 
@@ -24,6 +35,18 @@ using namespace std;
 
 namespace android::recoverymap {
 
+#define JPEGR_CHECK(x)          \
+  {                             \
+    status_t status = (x);      \
+    if ((status) != NO_ERROR) { \
+      return status;            \
+    }                           \
+  }
+
+// Map is quarter res / sixteenth size
+static const size_t kMapDimensionScaleFactor = 4;
+
+
 /*
  * Helper function used for generating XMP metadata.
  *
@@ -39,7 +62,7 @@ string Name(const string &prefix, const string &suffix) {
 
 status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
                                   jr_uncompressed_ptr uncompressed_yuv_420_image,
-                                  void* dest,
+                                  jr_compressed_ptr dest,
                                   int quality,
                                   jr_exif_ptr /* exif */,
                                   float /* hdr_ratio */) {
@@ -53,16 +76,44 @@ status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
     return ERROR_JPEGR_INVALID_INPUT_TYPE;
   }
 
-  // TBD
+  if (uncompressed_p010_image->width != uncompressed_yuv_420_image->width
+   || uncompressed_p010_image->height != uncompressed_yuv_420_image->height) {
+    return ERROR_JPEGR_RESOLUTION_MISMATCH;
+  }
+
+  jpegr_uncompressed_struct map;
+  JPEGR_CHECK(generateRecoveryMap(uncompressed_yuv_420_image, uncompressed_p010_image, &map));
+  std::unique_ptr<uint8_t[]> map_data;
+  map_data.reset(reinterpret_cast<uint8_t*>(map.data));
+
+  jpegr_compressed_struct compressed_map;
+  std::unique_ptr<uint8_t[]> compressed_map_data =
+      std::make_unique<uint8_t[]>(map.width * map.height);
+  compressed_map.data = compressed_map_data.get();
+  JPEGR_CHECK(compressRecoveryMap(&map, &compressed_map));
+
+  JpegEncoder jpeg_encoder;
+  // TODO: what quality to use?
+  // TODO: ICC data - need color space information
+  if (!jpeg_encoder.compressImage(uncompressed_yuv_420_image->data,
+                                  uncompressed_yuv_420_image->width,
+                                  uncompressed_yuv_420_image->height, 95, nullptr, 0)) {
+    return ERROR_JPEGR_ENCODE_ERROR;
+  }
+  jpegr_compressed_struct jpeg;
+  jpeg.data = jpeg_encoder.getCompressedImagePtr();
+  jpeg.length = jpeg_encoder.getCompressedImageSize();
+
+  JPEGR_CHECK(appendRecoveryMap(&jpeg, &compressed_map, dest));
+
   return NO_ERROR;
 }
 
 status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
                                   jr_uncompressed_ptr uncompressed_yuv_420_image,
-                                  void* compressed_jpeg_image,
-                                  void* dest,
+                                  jr_compressed_ptr compressed_jpeg_image,
+                                  jr_compressed_ptr dest,
                                   float /* hdr_ratio */) {
-
   if (uncompressed_p010_image == nullptr
    || uncompressed_yuv_420_image == nullptr
    || compressed_jpeg_image == nullptr
@@ -70,13 +121,30 @@ status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  // TBD
+  if (uncompressed_p010_image->width != uncompressed_yuv_420_image->width
+   || uncompressed_p010_image->height != uncompressed_yuv_420_image->height) {
+    return ERROR_JPEGR_RESOLUTION_MISMATCH;
+  }
+
+  jpegr_uncompressed_struct map;
+  JPEGR_CHECK(generateRecoveryMap(uncompressed_yuv_420_image, uncompressed_p010_image, &map));
+  std::unique_ptr<uint8_t[]> map_data;
+  map_data.reset(reinterpret_cast<uint8_t*>(map.data));
+
+  jpegr_compressed_struct compressed_map;
+  std::unique_ptr<uint8_t[]> compressed_map_data =
+      std::make_unique<uint8_t[]>(map.width * map.height);
+  compressed_map.data = compressed_map_data.get();
+  JPEGR_CHECK(compressRecoveryMap(&map, &compressed_map));
+
+  JPEGR_CHECK(appendRecoveryMap(compressed_jpeg_image, &compressed_map, dest));
+
   return NO_ERROR;
 }
 
 status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
-                                  void* compressed_jpeg_image,
-                                  void* dest,
+                                  jr_compressed_ptr compressed_jpeg_image,
+                                  jr_compressed_ptr dest,
                                   float /* hdr_ratio */) {
   if (uncompressed_p010_image == nullptr
    || compressed_jpeg_image == nullptr
@@ -84,11 +152,37 @@ status_t RecoveryMap::encodeJPEGR(jr_uncompressed_ptr uncompressed_p010_image,
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  // TBD
+  JpegDecoder jpeg_decoder;
+  if (!jpeg_decoder.decompressImage(compressed_jpeg_image->data, compressed_jpeg_image->length)) {
+    return ERROR_JPEGR_DECODE_ERROR;
+  }
+  jpegr_uncompressed_struct uncompressed_yuv_420_image;
+  uncompressed_yuv_420_image.data = jpeg_decoder.getDecompressedImagePtr();
+  uncompressed_yuv_420_image.width = jpeg_decoder.getDecompressedImageWidth();
+  uncompressed_yuv_420_image.height = jpeg_decoder.getDecompressedImageHeight();
+
+  if (uncompressed_p010_image->width != uncompressed_yuv_420_image.width
+   || uncompressed_p010_image->height != uncompressed_yuv_420_image.height) {
+    return ERROR_JPEGR_RESOLUTION_MISMATCH;
+  }
+
+  jpegr_uncompressed_struct map;
+  JPEGR_CHECK(generateRecoveryMap(&uncompressed_yuv_420_image, uncompressed_p010_image, &map));
+  std::unique_ptr<uint8_t[]> map_data;
+  map_data.reset(reinterpret_cast<uint8_t*>(map.data));
+
+  jpegr_compressed_struct compressed_map;
+  std::unique_ptr<uint8_t[]> compressed_map_data =
+      std::make_unique<uint8_t[]>(map.width * map.height);
+  compressed_map.data = compressed_map_data.get();
+  JPEGR_CHECK(compressRecoveryMap(&map, &compressed_map));
+
+  JPEGR_CHECK(appendRecoveryMap(compressed_jpeg_image, &compressed_map, dest));
+
   return NO_ERROR;
 }
 
-status_t RecoveryMap::decodeJPEGR(void* compressed_jpegr_image,
+status_t RecoveryMap::decodeJPEGR(jr_compressed_ptr compressed_jpegr_image,
                                   jr_uncompressed_ptr dest,
                                   jr_exif_ptr /* exif */,
                                   bool /* request_sdr */) {
@@ -96,27 +190,67 @@ status_t RecoveryMap::decodeJPEGR(void* compressed_jpegr_image,
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  // TBD
+  jpegr_compressed_struct compressed_map;
+  JPEGR_CHECK(extractRecoveryMap(compressed_jpegr_image, &compressed_map));
+
+  jpegr_uncompressed_struct map;
+  JPEGR_CHECK(decompressRecoveryMap(&compressed_map, &map));
+
+  JpegDecoder jpeg_decoder;
+  if (!jpeg_decoder.decompressImage(compressed_jpegr_image->data, compressed_jpegr_image->length)) {
+    return ERROR_JPEGR_DECODE_ERROR;
+  }
+
+  jpegr_uncompressed_struct uncompressed_yuv_420_image;
+  uncompressed_yuv_420_image.data = jpeg_decoder.getDecompressedImagePtr();
+  uncompressed_yuv_420_image.width = jpeg_decoder.getDecompressedImageWidth();
+  uncompressed_yuv_420_image.height = jpeg_decoder.getDecompressedImageHeight();
+
+  JPEGR_CHECK(applyRecoveryMap(&uncompressed_yuv_420_image, &map, dest));
+
   return NO_ERROR;
 }
 
-status_t RecoveryMap::decodeRecoveryMap(jr_compressed_ptr compressed_recovery_map,
-                                        jr_uncompressed_ptr dest) {
+status_t RecoveryMap::decompressRecoveryMap(jr_compressed_ptr compressed_recovery_map,
+                                            jr_uncompressed_ptr dest) {
   if (compressed_recovery_map == nullptr || dest == nullptr) {
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  // TBD
+  JpegDecoder jpeg_decoder;
+  if (!jpeg_decoder.decompressImage(compressed_recovery_map->data,
+                                    compressed_recovery_map->length)) {
+    return ERROR_JPEGR_DECODE_ERROR;
+  }
+
+  dest->data = jpeg_decoder.getDecompressedImagePtr();
+  dest->width = jpeg_decoder.getDecompressedImageWidth();
+  dest->height = jpeg_decoder.getDecompressedImageHeight();
+
   return NO_ERROR;
 }
 
-status_t RecoveryMap::encodeRecoveryMap(jr_uncompressed_ptr uncompressed_recovery_map,
-                                        jr_compressed_ptr dest) {
+status_t RecoveryMap::compressRecoveryMap(jr_uncompressed_ptr uncompressed_recovery_map,
+                                          jr_compressed_ptr dest) {
   if (uncompressed_recovery_map == nullptr || dest == nullptr) {
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  // TBD
+  // TODO: should we have ICC data?
+  JpegEncoder jpeg_encoder;
+  if (!jpeg_encoder.compressImage(uncompressed_recovery_map->data, uncompressed_recovery_map->width,
+                                  uncompressed_recovery_map->height, 85, nullptr, 0,
+                                  true /* isSingleChannel */)) {
+    return ERROR_JPEGR_ENCODE_ERROR;
+  }
+
+  if (dest->length < jpeg_encoder.getCompressedImageSize()) {
+    return ERROR_JPEGR_BUFFER_TOO_SMALL;
+  }
+
+  memcpy(dest->data, jpeg_encoder.getCompressedImagePtr(), jpeg_encoder.getCompressedImageSize());
+  dest->length = jpeg_encoder.getCompressedImageSize();
+
   return NO_ERROR;
 }
 
@@ -129,7 +263,51 @@ status_t RecoveryMap::generateRecoveryMap(jr_uncompressed_ptr uncompressed_yuv_4
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  // TBD
+  if (uncompressed_yuv_420_image->width != uncompressed_p010_image->width
+   || uncompressed_yuv_420_image->height != uncompressed_p010_image->height) {
+    return ERROR_JPEGR_RESOLUTION_MISMATCH;
+  }
+
+  size_t image_width = uncompressed_yuv_420_image->width;
+  size_t image_height = uncompressed_yuv_420_image->height;
+  size_t map_width = image_width / kMapDimensionScaleFactor;
+  size_t map_height = image_height / kMapDimensionScaleFactor;
+
+  dest->width = map_width;
+  dest->height = map_height;
+  dest->data = new uint8_t[map_width * map_height];
+  std::unique_ptr<uint8_t[]> map_data;
+  map_data.reset(reinterpret_cast<uint8_t*>(dest->data));
+
+  uint16_t yp_hdr_max = 0;
+  for (size_t y = 0; y < image_height; ++y) {
+    for (size_t x = 0; x < image_width; ++x) {
+      size_t pixel_idx =  x + y * image_width;
+      uint16_t yp_hdr = reinterpret_cast<uint8_t*>(uncompressed_yuv_420_image->data)[pixel_idx];
+      if (yp_hdr > yp_hdr_max) {
+        yp_hdr_max = yp_hdr;
+      }
+    }
+  }
+
+  float y_hdr_max_nits = hlgInvOetf(yp_hdr_max);
+  float hdr_ratio = y_hdr_max_nits / kSdrWhiteNits;
+
+  for (size_t y = 0; y < map_height; ++y) {
+    for (size_t x = 0; x < map_width; ++x) {
+      float yp_sdr = sampleYuv420Y(uncompressed_yuv_420_image, kMapDimensionScaleFactor, x, y);
+      float yp_hdr = sampleP010Y(uncompressed_p010_image, kMapDimensionScaleFactor, x, y);
+
+      float y_sdr_nits = srgbInvOetf(yp_sdr);
+      float y_hdr_nits = hlgInvOetf(yp_hdr);
+
+      size_t pixel_idx =  x + y * map_width;
+      reinterpret_cast<uint8_t*>(dest->data)[pixel_idx] =
+          encodeRecovery(y_sdr_nits, y_hdr_nits, hdr_ratio);
+    }
+  }
+
+  map_data.release();
   return NO_ERROR;
 }
 
@@ -142,11 +320,44 @@ status_t RecoveryMap::applyRecoveryMap(jr_uncompressed_ptr uncompressed_yuv_420_
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
 
-  // TBD
+  // TODO: need to get this from the XMP; should probably be a function
+  // parameter
+  float hdr_ratio = 4.0f;
+
+  size_t width = uncompressed_yuv_420_image->width;
+  size_t height = uncompressed_yuv_420_image->height;
+
+  dest->width = width;
+  dest->height = height;
+  size_t pixel_count = width * height;
+
+  for (size_t y = 0; y < height; ++y) {
+    for (size_t x = 0; x < width; ++x) {
+      size_t pixel_y_idx =  x + y * width;
+
+      size_t pixel_uv_idx = x / 2 + (y / 2) * (width / 2);
+
+      Color ypuv_sdr = getYuv420Pixel(uncompressed_yuv_420_image, x, y);
+      Color rgbp_sdr = srgbYuvToRgb(ypuv_sdr);
+      Color rgb_sdr = srgbInvOetf(rgbp_sdr);
+
+      float recovery = sampleMap(uncompressed_recovery_map, kMapDimensionScaleFactor, x, y);
+      Color rgb_hdr = applyRecovery(rgb_sdr, recovery, hdr_ratio);
+
+      Color rgbp_hdr = hlgOetf(rgb_hdr);
+      Color ypuv_hdr = bt2100RgbToYuv(rgbp_hdr);
+
+      reinterpret_cast<uint16_t*>(dest->data)[pixel_y_idx] = ypuv_hdr.r;
+      reinterpret_cast<uint16_t*>(dest->data)[pixel_count + pixel_uv_idx] = ypuv_hdr.g;
+      reinterpret_cast<uint16_t*>(dest->data)[pixel_count + pixel_uv_idx + 1] = ypuv_hdr.b;
+    }
+  }
+
   return NO_ERROR;
 }
 
-status_t RecoveryMap::extractRecoveryMap(void* compressed_jpegr_image, jr_compressed_ptr dest) {
+status_t RecoveryMap::extractRecoveryMap(jr_compressed_ptr compressed_jpegr_image,
+                                         jr_compressed_ptr dest) {
   if (compressed_jpegr_image == nullptr || dest == nullptr) {
     return ERROR_JPEGR_INVALID_NULL_PTR;
   }
@@ -155,9 +366,9 @@ status_t RecoveryMap::extractRecoveryMap(void* compressed_jpegr_image, jr_compre
   return NO_ERROR;
 }
 
-status_t RecoveryMap::appendRecoveryMap(void* compressed_jpeg_image,
-                                     jr_compressed_ptr compressed_recovery_map,
-                                     void* dest) {
+status_t RecoveryMap::appendRecoveryMap(jr_compressed_ptr compressed_jpeg_image,
+                                        jr_compressed_ptr compressed_recovery_map,
+                                        jr_compressed_ptr dest) {
   if (compressed_jpeg_image == nullptr
    || compressed_recovery_map == nullptr
    || dest == nullptr) {
diff --git a/libs/jpegrecoverymap/recoverymapmath.cpp b/libs/jpegrecoverymap/recoverymapmath.cpp
new file mode 100644
index 0000000000..3e110bcd26
--- /dev/null
+++ b/libs/jpegrecoverymap/recoverymapmath.cpp
@@ -0,0 +1,169 @@
+/*
+ * 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.
+ */
+
+#include <cmath>
+
+#include <jpegrecoverymap/recoverymapmath.h>
+
+namespace android::recoverymap {
+
+static const float kBt2100R = 0.2627f, kBt2100G = 0.6780f, kBt2100B = 0.0593f;
+static const float kBt2100Cb = 1.8814f, kBt2100Cr = 1.4746f;
+
+Color bt2100RgbToYuv(Color e) {
+  float yp = kBt2100R * e.r + kBt2100G * e.g + kBt2100B * e.b;
+  return {{{yp, (e.b - yp) / kBt2100Cb, (e.r - yp) / kBt2100Cr }}};
+}
+
+static const float kSrgbRCr = 1.402f, kSrgbGCb = 0.34414f, kSrgbGCr = 0.71414f, kSrgbBCb = 1.772f;
+
+Color srgbYuvToRgb(Color e) {
+  return {{{ e.y + kSrgbRCr * e.v, e.y - kSrgbGCb * e.u - kSrgbGCr * e.v, e.y + kSrgbBCb * e.u }}};
+}
+
+float srgbInvOetf(float e) {
+  if (e <= 0.04045f) {
+    return e / 12.92f;
+  } else {
+    return pow((e + 0.055f) / 1.055f, 2.4);
+  }
+}
+
+Color srgbInvOetf(Color e) {
+  return {{{ srgbInvOetf(e.r), srgbInvOetf(e.g), srgbInvOetf(e.b) }}};
+}
+
+static const float kHlgA = 0.17883277f, kHlgB = 0.28466892f, kHlgC = 0.55991073;
+
+float hlgInvOetf(float e) {
+  if (e <= 0.5f) {
+    return pow(e, 2.0f) / 3.0f;
+  } else {
+    return (exp((e - kHlgC) / kHlgA) + kHlgB) / 12.0f;
+  }
+}
+
+float hlgOetf(float e) {
+  if (e <= 1.0f/12.0f) {
+    return sqrt(3.0f * e);
+  } else {
+    return kHlgA * log(12.0f * e - kHlgB) + kHlgC;
+  }
+}
+
+Color hlgOetf(Color e) {
+  return {{{ hlgOetf(e.r), hlgOetf(e.g), hlgOetf(e.b) }}};
+}
+
+uint8_t EncodeRecovery(float y_sdr, float y_hdr, float hdr_ratio) {
+  float gain = 1.0f;
+  if (y_sdr > 0.0f) {
+    gain = y_hdr / y_sdr;
+  }
+
+  if (gain < -hdr_ratio) gain = -hdr_ratio;
+  if (gain > hdr_ratio) gain = hdr_ratio;
+
+  return static_cast<uint8_t>(log2(gain) / log2(hdr_ratio) * 127.5f  + 127.5f);
+}
+
+float applyRecovery(float y_sdr, float recovery, float hdr_ratio) {
+  return exp2(log2(y_sdr) + recovery * log2(hdr_ratio));
+}
+
+// TODO: do we need something more clever for filtering either the map or images
+// to generate the map?
+
+static float mapUintToFloat(uint8_t map_uint) {
+  return (static_cast<float>(map_uint) - 127.5f) / 127.5f;
+}
+
+float sampleMap(jr_uncompressed_ptr map, size_t map_scale_factor, size_t x, size_t y) {
+  float x_map = static_cast<float>(x) / static_cast<float>(map_scale_factor);
+  float y_map = static_cast<float>(y) / static_cast<float>(map_scale_factor);
+
+  size_t x_lower = static_cast<size_t>(floor(x_map));
+  size_t x_upper = x_lower + 1;
+  size_t y_lower = static_cast<size_t>(floor(y_map));
+  size_t y_upper = y_lower + 1;
+
+  float x_influence = x_map - static_cast<float>(x_lower);
+  float y_influence = y_map - static_cast<float>(y_lower);
+
+  float e1 = mapUintToFloat(reinterpret_cast<uint8_t*>(map->data)[x_lower + y_lower * map->width]);
+  float e2 = mapUintToFloat(reinterpret_cast<uint8_t*>(map->data)[x_lower + y_upper * map->width]);
+  float e3 = mapUintToFloat(reinterpret_cast<uint8_t*>(map->data)[x_upper + y_lower * map->width]);
+  float e4 = mapUintToFloat(reinterpret_cast<uint8_t*>(map->data)[x_upper + y_upper * map->width]);
+
+  return e1 * (x_influence + y_influence) / 2.0f
+      + e2 * (x_influence + 1.0f - y_influence) / 2.0f
+      + e3 * (1.0f - x_influence + y_influence) / 2.0f
+      + e4 * (1.0f - x_influence + 1.0f - y_influence) / 2.0f;
+}
+
+Color getYuv420Pixel(jr_uncompressed_ptr image, size_t x, size_t y) {
+  size_t pixel_count = image->width * image->height;
+
+  size_t pixel_y_idx = x + y * image->width;
+  size_t pixel_uv_idx = x / 2 + (y / 2) * (image->width / 2);
+
+  uint8_t y_uint = reinterpret_cast<uint8_t*>(image->data)[pixel_y_idx];
+  uint8_t u_uint = reinterpret_cast<uint8_t*>(image->data)[pixel_count + pixel_uv_idx];
+  uint8_t v_uint = reinterpret_cast<uint8_t*>(image->data)[pixel_count * 5 / 4 + pixel_uv_idx];
+
+  // 128 bias for UV given we are using jpeglib; see:
+  // https://github.com/kornelski/libjpeg/blob/master/structure.doc
+  return {{{ static_cast<float>(y_uint) / 255.0f,
+             (static_cast<float>(u_uint) - 128.0f) / 255.0f,
+             (static_cast<float>(v_uint) - 128.0f) / 255.0f }}};
+}
+
+typedef float (*sampleComponentFn)(jr_uncompressed_ptr, size_t, size_t);
+
+static float sampleComponent(jr_uncompressed_ptr image, size_t map_scale_factor, size_t x, size_t y,
+                             sampleComponentFn sample_fn) {
+  float e = 0.0f;
+  for (size_t dy = 0; dy < map_scale_factor; ++dy) {
+    for (size_t dx = 0; dx < map_scale_factor; ++dx) {
+      e += sample_fn(image, x * map_scale_factor + dx, y * map_scale_factor + dy);
+    }
+  }
+
+  return e / static_cast<float>(map_scale_factor * map_scale_factor);
+}
+
+static float getYuv420Y(jr_uncompressed_ptr image, size_t x, size_t y) {
+  size_t pixel_idx = x + y * image->width;
+  uint8_t y_uint = reinterpret_cast<uint8_t*>(image->data)[pixel_idx];
+  return static_cast<float>(y_uint) / 255.0f;
+}
+
+
+float sampleYuv420Y(jr_uncompressed_ptr image, size_t map_scale_factor, size_t x, size_t y) {
+  return sampleComponent(image, map_scale_factor, x, y, getYuv420Y);
+}
+
+static float getP010Y(jr_uncompressed_ptr image, size_t x, size_t y) {
+  size_t pixel_idx = x + y * image->width;
+  uint8_t y_uint = reinterpret_cast<uint16_t*>(image->data)[pixel_idx];
+  // Expecting narrow range input
+  return (static_cast<float>(y_uint) - 64.0f) / 960.0f;
+}
+
+float sampleP010Y(jr_uncompressed_ptr image, size_t map_scale_factor, size_t x, size_t y) {
+  return sampleComponent(image, map_scale_factor, x, y, getP010Y);
+}
+} // namespace android::recoverymap
diff --git a/libs/jpegrecoverymap/tests/Android.bp b/libs/jpegrecoverymap/tests/Android.bp
index 7f37f611c7..41af991093 100644
--- a/libs/jpegrecoverymap/tests/Android.bp
+++ b/libs/jpegrecoverymap/tests/Android.bp
@@ -62,4 +62,4 @@ cc_test {
         "libjpegdecoder",
         "libgtest",
     ],
-}
\ No newline at end of file
+}