Updates to jpeg encoder helper class

- Add support for more sampling formats
- RGB jpeg encode now defaults to 444 sampling format instead of 420
- For unaligned blocks memset chroma planes to 128 to create black borders
- Add more fail safe checks, diagnostics
- Ensure memory is released during encoding failures

Test: ./ultrahdr_unit_test
Test: ./ultrahdr_enc_fuzzer

2 files changed, 213 insertions, 282 deletions

diff --git a/lib/src/jpegencoderhelper.cpp b/lib/src/jpegencoderhelper.cpp
index 7097fef..9db6a3c 100644
--- a/lib/src/jpegencoderhelper.cpp
+++ b/lib/src/jpegencoderhelper.cpp
@@ -14,7 +14,12 @@
  * limitations under the License.
  */
 
+#include <errno.h>
+#include <setjmp.h>
+
+#include <cmath>
 #include <cstring>
+#include <map>
 #include <memory>
 #include <string>
 
@@ -24,68 +29,75 @@
 
 namespace ultrahdr {
 
-// The destination manager that can access |mResultBuffer| in JpegEncoderHelper.
-struct destination_mgr {
-  struct jpeg_destination_mgr mgr;
-  JpegEncoderHelper* encoder;
+/*!\brief map of sub sampling format and jpeg h_samp_factor, v_samp_factor */
+std::map<JpegEncoderHelper::jpg_inp_fmt_t, std::vector<int>> sample_factors = {
+    {JpegEncoderHelper::GRAYSCALE,
+     {1 /*h0*/, 1 /*v0*/, 0 /*h1*/, 0 /*v1*/, 0 /*h2*/, 0 /*v2*/, 1 /*maxh*/, 1 /*maxv*/}},
+    {JpegEncoderHelper::YUV444,
+     {1 /*h0*/, 1 /*v0*/, 1 /*h1*/, 1 /*v1*/, 1 /*h2*/, 1 /*v2*/, 1 /*maxh*/, 1 /*maxv*/}},
+    {JpegEncoderHelper::YUV440,
+     {1 /*h0*/, 2 /*v0*/, 1 /*h1*/, 1 /*v1*/, 1 /*h2*/, 1 /*v2*/, 1 /*maxh*/, 2 /*maxv*/}},
+    {JpegEncoderHelper::YUV422,
+     {2 /*h0*/, 1 /*v0*/, 1 /*h1*/, 1 /*v1*/, 1 /*h2*/, 1 /*v2*/, 2 /*maxh*/, 1 /*maxv*/}},
+    {JpegEncoderHelper::YUV420,
+     {2 /*h0*/, 2 /*v0*/, 1 /*h1*/, 1 /*v1*/, 1 /*h2*/, 1 /*v2*/, 2 /*maxh*/, 2 /*maxv*/}},
+    {JpegEncoderHelper::YUV411,
+     {4 /*h0*/, 1 /*v0*/, 1 /*h1*/, 1 /*v1*/, 1 /*h2*/, 1 /*v2*/, 4 /*maxh*/, 1 /*maxv*/}},
+    {JpegEncoderHelper::YUV410,
+     {4 /*h0*/, 2 /*v0*/, 1 /*h1*/, 1 /*v1*/, 1 /*h2*/, 1 /*v2*/, 4 /*maxh*/, 2 /*maxv*/}},
+    {JpegEncoderHelper::RGB,
+     {1 /*h0*/, 1 /*v0*/, 1 /*h1*/, 1 /*v1*/, 1 /*h2*/, 1 /*v2*/, 1 /*maxh*/, 1 /*maxv*/}},
 };
 
-JpegEncoderHelper::JpegEncoderHelper() {}
-
-JpegEncoderHelper::~JpegEncoderHelper() {}
-
-bool JpegEncoderHelper::compressImage(const uint8_t* yBuffer, const uint8_t* uvBuffer, int width,
-                                      int height, int lumaStride, int chromaStride, int quality,
-                                      const void* iccBuffer, unsigned int iccSize) {
-  mResultBuffer.clear();
-  if (!encode(yBuffer, uvBuffer, width, height, lumaStride, chromaStride, quality, iccBuffer,
-              iccSize)) {
-    return false;
-  }
-  ALOGV("Compressed JPEG: %d[%dx%d] -> %zu bytes", (width * height * 12) / 8, width, height,
-        mResultBuffer.size());
-  return true;
+/*!\brief jpeg encoder library destination manager callback functions implementation */
+
+/*!\brief  called by jpeg_start_compress() before any data is actually written. This function is
+ * expected to initialize fields next_output_byte (place to write encoded output) and
+ * free_in_buffer (size of the buffer supplied) of jpeg destination manager. free_in_buffer must
+ * be initialized to a positive value.*/
+static void initDestination(j_compress_ptr cinfo) {
+  destination_mgr_impl* dest = reinterpret_cast<destination_mgr_impl*>(cinfo->dest);
+  std::vector<JOCTET>& buffer = dest->mResultBuffer;
+  buffer.resize(dest->kBlockSize);
+  dest->next_output_byte = &buffer[0];
+  dest->free_in_buffer = buffer.size();
 }
 
-bool JpegEncoderHelper::compressImage(const uint8_t* buffer, int width, int height, int quality,
-                                      const void* iccBuffer, unsigned int iccSize) {
-  mResultBuffer.clear();
-  if (!encode(buffer, width, height, quality, iccBuffer, iccSize)) {
-    return false;
-  }
-  ALOGV("Compressed JPEG: [%dx%d] -> %zu bytes", width, height, mResultBuffer.size());
+/*!\brief  called if buffer provided for storing encoded data is exhausted during encoding. This
+ * function is expected to consume the encoded output and provide fresh buffer to continue
+ * encoding. */
+static boolean emptyOutputBuffer(j_compress_ptr cinfo) {
+  destination_mgr_impl* dest = reinterpret_cast<destination_mgr_impl*>(cinfo->dest);
+  std::vector<JOCTET>& buffer = dest->mResultBuffer;
+  size_t oldsize = buffer.size();
+  buffer.resize(oldsize + dest->kBlockSize);
+  dest->next_output_byte = &buffer[oldsize];
+  dest->free_in_buffer = dest->kBlockSize;
   return true;
 }
 
-void* JpegEncoderHelper::getCompressedImagePtr() { return mResultBuffer.data(); }
-
-size_t JpegEncoderHelper::getCompressedImageSize() { return mResultBuffer.size(); }
-
-void JpegEncoderHelper::initDestination(j_compress_ptr cinfo) {
-  destination_mgr* dest = reinterpret_cast<destination_mgr*>(cinfo->dest);
-  std::vector<JOCTET>& buffer = dest->encoder->mResultBuffer;
-  buffer.resize(kBlockSize);
-  dest->mgr.next_output_byte = &buffer[0];
-  dest->mgr.free_in_buffer = buffer.size();
+/*!\brief  called by jpeg_finish_compress() to flush out all the remaining encoded data. client
+ * can use either next_output_byte or free_in_buffer to determine how much data is in the buffer.
+ */
+static void terminateDestination(j_compress_ptr cinfo) {
+  destination_mgr_impl* dest = reinterpret_cast<destination_mgr_impl*>(cinfo->dest);
+  std::vector<JOCTET>& buffer = dest->mResultBuffer;
+  buffer.resize(buffer.size() - dest->free_in_buffer);
 }
 
-boolean JpegEncoderHelper::emptyOutputBuffer(j_compress_ptr cinfo) {
-  destination_mgr* dest = reinterpret_cast<destination_mgr*>(cinfo->dest);
-  std::vector<JOCTET>& buffer = dest->encoder->mResultBuffer;
-  size_t oldsize = buffer.size();
-  buffer.resize(oldsize + kBlockSize);
-  dest->mgr.next_output_byte = &buffer[oldsize];
-  dest->mgr.free_in_buffer = kBlockSize;
-  return true;
-}
+/*!\brief module for managing error */
+struct jpeg_error_mgr_impl : jpeg_error_mgr {
+  jmp_buf setjmp_buffer;
+};
 
-void JpegEncoderHelper::terminateDestination(j_compress_ptr cinfo) {
-  destination_mgr* dest = reinterpret_cast<destination_mgr*>(cinfo->dest);
-  std::vector<JOCTET>& buffer = dest->encoder->mResultBuffer;
-  buffer.resize(buffer.size() - dest->mgr.free_in_buffer);
+/*!\brief jpeg encoder library error manager callback function implementations */
+static void jpegrerror_exit(j_common_ptr cinfo) {
+  jpeg_error_mgr_impl* err = reinterpret_cast<jpeg_error_mgr_impl*>(cinfo->err);
+  longjmp(err->setjmp_buffer, 1);
 }
 
-void JpegEncoderHelper::outputErrorMessage(j_common_ptr cinfo) {
+/* receive most recent jpeg error message and print */
+static void outputErrorMessage(j_common_ptr cinfo) {
   char buffer[JMSG_LENGTH_MAX];
 
   /* Create the message */
@@ -93,243 +105,153 @@ void JpegEncoderHelper::outputErrorMessage(j_common_ptr cinfo) {
   ALOGE("%s\n", buffer);
 }
 
-bool JpegEncoderHelper::encode(const uint8_t* yBuffer, const uint8_t* uvBuffer, int width,
-                               int height, int lumaStride, int chromaStride, int quality,
-                               const void* iccBuffer, unsigned int iccSize) {
-  jpeg_compress_struct cinfo;
-  jpeg_error_mgr jerr;
-
-  cinfo.err = jpeg_std_error(&jerr);
-  cinfo.err->output_message = &outputErrorMessage;
-  jpeg_create_compress(&cinfo);
-  setJpegDestination(&cinfo);
-  setJpegCompressStruct(width, height, quality, &cinfo,
-                        uvBuffer == nullptr ? ENCODE_TO_SINGLE_CHANNEL : ENCODE_TO_YCBCR);
-  jpeg_start_compress(&cinfo, TRUE);
-  if (iccBuffer != nullptr && iccSize > 0) {
-    jpeg_write_marker(&cinfo, JPEG_APP0 + 2, static_cast<const JOCTET*>(iccBuffer), iccSize);
-  }
-  bool status = cinfo.num_components == 1
-                    ? compressY(&cinfo, yBuffer, lumaStride)
-                    : compressYuv(&cinfo, yBuffer, uvBuffer, lumaStride, chromaStride);
-  jpeg_finish_compress(&cinfo);
-  jpeg_destroy_compress(&cinfo);
-
-  return status;
+bool JpegEncoderHelper::compressImage(const uint8_t* planes[3], const size_t strides[3],
+                                      const int width, const int height, const jpg_inp_fmt_t format,
+                                      const int qfactor, const void* iccBuffer,
+                                      const unsigned int iccSize) {
+  return encode(planes, strides, width, height, format, qfactor, iccBuffer, iccSize);
 }
 
-bool JpegEncoderHelper::encode(const uint8_t* buffer, int width, int height, int quality,
-                               const void* iccBuffer, unsigned int iccSize) {
+bool JpegEncoderHelper::encode(const uint8_t* planes[3], const size_t strides[3], const int width,
+                               const int height, const jpg_inp_fmt_t format, const int qfactor,
+                               const void* iccBuffer, const unsigned int iccSize) {
   jpeg_compress_struct cinfo;
-  jpeg_error_mgr jerr;
-
-  cinfo.err = jpeg_std_error(&jerr);
-  cinfo.err->output_message = &outputErrorMessage;
-  jpeg_create_compress(&cinfo);
-  setJpegDestination(&cinfo);
-  setJpegCompressStruct(width, height, quality, &cinfo, ENCODE_TO_RGB);
-  jpeg_start_compress(&cinfo, TRUE);
-  if (iccBuffer != nullptr && iccSize > 0) {
-    jpeg_write_marker(&cinfo, JPEG_APP0 + 2, static_cast<const JOCTET*>(iccBuffer), iccSize);
-  }
-
-  while (cinfo.next_scanline < cinfo.image_height) {
-    JSAMPROW row_pointer[1];
-    row_pointer[0] = const_cast<JSAMPROW>(&buffer[cinfo.next_scanline * width * 3]);
-    (void)jpeg_write_scanlines(&cinfo, row_pointer, 1);
-  }
-
-  jpeg_finish_compress(&cinfo);
-  jpeg_destroy_compress(&cinfo);
+  jpeg_error_mgr_impl myerr;
 
-  return true;
-}
-
-void JpegEncoderHelper::setJpegDestination(jpeg_compress_struct* cinfo) {
-  destination_mgr* dest = static_cast<struct destination_mgr*>(
-      (*cinfo->mem->alloc_small)((j_common_ptr)cinfo, JPOOL_PERMANENT, sizeof(destination_mgr)));
-  dest->encoder = this;
-  dest->mgr.init_destination = &initDestination;
-  dest->mgr.empty_output_buffer = &emptyOutputBuffer;
-  dest->mgr.term_destination = &terminateDestination;
-  cinfo->dest = reinterpret_cast<struct jpeg_destination_mgr*>(dest);
-}
-
-void JpegEncoderHelper::setJpegCompressStruct(int width, int height, int quality,
-                                              jpeg_compress_struct* cinfo,
-                                              encode_mode_t encodeMode) {
-  cinfo->image_width = width;
-  cinfo->image_height = height;
-
-  if (encodeMode == ENCODE_TO_SINGLE_CHANNEL) {
-    cinfo->input_components = 1;
-    cinfo->in_color_space = JCS_GRAYSCALE;
-  } else if (encodeMode == ENCODE_TO_YCBCR) {
-    cinfo->input_components = 3;
-    cinfo->in_color_space = JCS_YCbCr;
-  } else {
-    cinfo->input_components = 3;
-    cinfo->in_color_space = JCS_RGB;
+  if (sample_factors.find(format) == sample_factors.end()) {
+    ALOGE("unrecognized format %d", format);
+    return false;
   }
-
-  jpeg_set_defaults(cinfo);
-  jpeg_set_quality(cinfo, quality, TRUE);
-
-  if (encodeMode == ENCODE_TO_SINGLE_CHANNEL || encodeMode == ENCODE_TO_YCBCR) {
-    cinfo->raw_data_in = TRUE;
-    cinfo->dct_method = JDCT_ISLOW;
-    cinfo->comp_info[0].h_samp_factor = cinfo->in_color_space == JCS_GRAYSCALE ? 1 : 2;
-    cinfo->comp_info[0].v_samp_factor = cinfo->in_color_space == JCS_GRAYSCALE ? 1 : 2;
-    for (int i = 1; i < cinfo->num_components; i++) {
-      cinfo->comp_info[i].h_samp_factor = 1;
-      cinfo->comp_info[i].v_samp_factor = 1;
+  std::vector<int>& factors = sample_factors.find(format)->second;
+
+  cinfo.err = jpeg_std_error(&myerr);
+  myerr.error_exit = jpegrerror_exit;
+  myerr.output_message = outputErrorMessage;
+
+  if (0 == setjmp(myerr.setjmp_buffer)) {
+    jpeg_create_compress(&cinfo);
+
+    // initialize destination manager
+    mDestMgr.init_destination = &initDestination;
+    mDestMgr.empty_output_buffer = &emptyOutputBuffer;
+    mDestMgr.term_destination = &terminateDestination;
+    mDestMgr.mResultBuffer.clear();
+    cinfo.dest = reinterpret_cast<struct jpeg_destination_mgr*>(&mDestMgr);
+
+    // initialize confiurations parameters
+    cinfo.image_width = width;
+    cinfo.image_height = height;
+    if (format == RGB) {
+      cinfo.input_components = 3;
+      cinfo.in_color_space = JCS_RGB;
+    } else {
+      if (format == GRAYSCALE) {
+        cinfo.input_components = 1;
+        cinfo.in_color_space = JCS_GRAYSCALE;
+      } else {
+        cinfo.input_components = 3;
+        cinfo.in_color_space = JCS_YCbCr;
+      }
     }
-  }
-}
-
-bool JpegEncoderHelper::compressYuv(jpeg_compress_struct* cinfo, const uint8_t* yBuffer,
-                                    const uint8_t* uvBuffer, int lumaStride, int chromaStride) {
-  size_t chroma_plane_size = chromaStride * cinfo->image_height / 2;
-  uint8_t* y_plane = const_cast<uint8_t*>(yBuffer);
-  uint8_t* u_plane = const_cast<uint8_t*>(uvBuffer);
-  uint8_t* v_plane = const_cast<uint8_t*>(u_plane + chroma_plane_size);
-
-  const int aligned_width = ALIGNM(cinfo->image_width, kCompressBatchSize);
-  const bool need_luma_padding = (lumaStride < aligned_width);
-  const int aligned_chroma_width = ALIGNM(cinfo->image_width / 2, kCompressBatchSize / 2);
-  const bool need_chroma_padding = (chromaStride < aligned_chroma_width);
-
-  std::unique_ptr<uint8_t[]> empty = nullptr;
-  std::unique_ptr<uint8_t[]> y_mcu_row = nullptr;
-  std::unique_ptr<uint8_t[]> cb_mcu_row = nullptr;
-  std::unique_ptr<uint8_t[]> cr_mcu_row = nullptr;
-  uint8_t* y_mcu_row_ptr = nullptr;
-  uint8_t* cb_mcu_row_ptr = nullptr;
-  uint8_t* cr_mcu_row_ptr = nullptr;
-
-  JSAMPROW y[kCompressBatchSize];
-  JSAMPROW cb[kCompressBatchSize / 2];
-  JSAMPROW cr[kCompressBatchSize / 2];
-  JSAMPARRAY planes[3]{y, cb, cr};
-
-  if (cinfo->image_height % kCompressBatchSize != 0) {
-    empty = std::make_unique<uint8_t[]>(aligned_width);
-    memset(empty.get(), 0, aligned_width);
-  }
-
-  if (need_luma_padding) {
-    size_t mcu_row_size = aligned_width * kCompressBatchSize;
-    y_mcu_row = std::make_unique<uint8_t[]>(mcu_row_size);
-    y_mcu_row_ptr = y_mcu_row.get();
-    uint8_t* tmp = y_mcu_row_ptr;
-    for (int i = 0; i < kCompressBatchSize; ++i, tmp += aligned_width) {
-      memset(tmp + cinfo->image_width, 0, aligned_width - cinfo->image_width);
+    jpeg_set_defaults(&cinfo);
+    jpeg_set_quality(&cinfo, qfactor, TRUE);
+    for (int i = 0; i < cinfo.num_components; i++) {
+      cinfo.comp_info[i].h_samp_factor = factors[i * 2];
+      cinfo.comp_info[i].v_samp_factor = factors[i * 2 + 1];
+      mPlaneWidth[i] =
+          std::ceil(((float)cinfo.image_width * cinfo.comp_info[i].h_samp_factor) / factors[6]);
+      mPlaneHeight[i] =
+          std::ceil(((float)cinfo.image_height * cinfo.comp_info[i].v_samp_factor) / factors[7]);
     }
-  }
+    if (format != RGB) cinfo.raw_data_in = TRUE;
+    cinfo.dct_method = JDCT_ISLOW;
 
-  if (need_chroma_padding) {
-    size_t mcu_row_size = aligned_chroma_width * kCompressBatchSize / 2;
-    cb_mcu_row = std::make_unique<uint8_t[]>(mcu_row_size);
-    cb_mcu_row_ptr = cb_mcu_row.get();
-    cr_mcu_row = std::make_unique<uint8_t[]>(mcu_row_size);
-    cr_mcu_row_ptr = cr_mcu_row.get();
-    uint8_t* tmp1 = cb_mcu_row_ptr;
-    uint8_t* tmp2 = cr_mcu_row_ptr;
-    for (int i = 0; i < kCompressBatchSize / 2;
-         ++i, tmp1 += aligned_chroma_width, tmp2 += aligned_chroma_width) {
-      memset(tmp1 + cinfo->image_width / 2, 0, aligned_chroma_width - (cinfo->image_width / 2));
-      memset(tmp2 + cinfo->image_width / 2, 0, aligned_chroma_width - (cinfo->image_width / 2));
+    // start compress
+    jpeg_start_compress(&cinfo, TRUE);
+    if (iccBuffer != nullptr && iccSize > 0) {
+      jpeg_write_marker(&cinfo, JPEG_APP0 + 2, static_cast<const JOCTET*>(iccBuffer), iccSize);
     }
-  }
-
-  while (cinfo->next_scanline < cinfo->image_height) {
-    for (int i = 0; i < kCompressBatchSize; ++i) {
-      size_t scanline = cinfo->next_scanline + i;
-      if (scanline < cinfo->image_height) {
-        y[i] = y_plane + scanline * lumaStride;
-        if (need_luma_padding) {
-          uint8_t* tmp = y_mcu_row_ptr + i * aligned_width;
-          memcpy(tmp, y[i], cinfo->image_width);
-          y[i] = tmp;
+    if (format == RGB) {
+      while (cinfo.next_scanline < cinfo.image_height) {
+        JSAMPROW row_pointer[]{const_cast<JSAMPROW>(&planes[0][cinfo.next_scanline * strides[0]])};
+        JDIMENSION processed = jpeg_write_scanlines(&cinfo, row_pointer, 1);
+        if (1 != processed) {
+          ALOGE("jpeg_read_scanlines returned %d, expected %d", processed, 1);
+          jpeg_destroy_compress(&cinfo);
+          return false;
         }
-      } else {
-        y[i] = empty.get();
       }
-    }
-    // cb, cr only have half scanlines
-    for (int i = 0; i < kCompressBatchSize / 2; ++i) {
-      size_t scanline = cinfo->next_scanline / 2 + i;
-      if (scanline < cinfo->image_height / 2) {
-        int offset = scanline * chromaStride;
-        cb[i] = u_plane + offset;
-        cr[i] = v_plane + offset;
-        if (need_chroma_padding) {
-          uint8_t* tmp = cb_mcu_row_ptr + i * aligned_chroma_width;
-          memcpy(tmp, cb[i], cinfo->image_width / 2);
-          cb[i] = tmp;
-          tmp = cr_mcu_row_ptr + i * aligned_chroma_width;
-          memcpy(tmp, cr[i], cinfo->image_width / 2);
-          cr[i] = tmp;
-        }
-      } else {
-        cb[i] = cr[i] = empty.get();
+    } else {
+      if (!compressYCbCr(&cinfo, planes, strides)) {
+        jpeg_destroy_compress(&cinfo);
+        return false;
       }
     }
-    int processed = jpeg_write_raw_data(cinfo, planes, kCompressBatchSize);
-    if (processed != kCompressBatchSize) {
-      ALOGE("Number of processed lines does not equal input lines.");
-      return false;
-    }
+  } else {
+    cinfo.err->output_message((j_common_ptr)&cinfo);
+    jpeg_destroy_compress(&cinfo);
+    return false;
   }
+
+  jpeg_finish_compress(&cinfo);
+  jpeg_destroy_compress(&cinfo);
   return true;
 }
 
-bool JpegEncoderHelper::compressY(jpeg_compress_struct* cinfo, const uint8_t* yBuffer,
-                                  int lumaStride) {
-  uint8_t* y_plane = const_cast<uint8_t*>(yBuffer);
-
-  const int aligned_luma_width = ALIGNM(cinfo->image_width, kCompressBatchSize);
-  const bool need_luma_padding = (lumaStride < aligned_luma_width);
-
-  std::unique_ptr<uint8_t[]> empty = nullptr;
-  std::unique_ptr<uint8_t[]> y_mcu_row = nullptr;
-  uint8_t* y_mcu_row_ptr = nullptr;
-
-  JSAMPROW y[kCompressBatchSize];
-  JSAMPARRAY planes[1]{y};
-
-  if (cinfo->image_height % kCompressBatchSize != 0) {
-    empty = std::make_unique<uint8_t[]>(aligned_luma_width);
-    memset(empty.get(), 0, aligned_luma_width);
-  }
-
-  if (need_luma_padding) {
-    size_t mcu_row_size = aligned_luma_width * kCompressBatchSize;
-    y_mcu_row = std::make_unique<uint8_t[]>(mcu_row_size);
-    y_mcu_row_ptr = y_mcu_row.get();
-    uint8_t* tmp = y_mcu_row_ptr;
-    for (int i = 0; i < kCompressBatchSize; ++i, tmp += aligned_luma_width) {
-      memset(tmp + cinfo->image_width, 0, aligned_luma_width - cinfo->image_width);
+bool JpegEncoderHelper::compressYCbCr(jpeg_compress_struct* cinfo, const uint8_t* planes[3],
+                                      const size_t strides[3]) {
+  JSAMPROW mcuRows[kMaxNumComponents][2 * DCTSIZE];
+  JSAMPROW mcuRowsTmp[kMaxNumComponents][2 * DCTSIZE];
+  size_t alignedPlaneWidth[kMaxNumComponents]{};
+  JSAMPARRAY subImage[kMaxNumComponents];
+
+  for (int i = 0; i < cinfo->num_components; i++) {
+    alignedPlaneWidth[i] = ALIGNM(mPlaneWidth[i], DCTSIZE);
+    if (strides[i] < alignedPlaneWidth[i]) {
+      mPlanesMCURow[i] = std::make_unique<uint8_t[]>(alignedPlaneWidth[i] * DCTSIZE *
+                                                     cinfo->comp_info[i].v_samp_factor);
+      uint8_t* mem = mPlanesMCURow[i].get();
+      for (int j = 0; j < DCTSIZE * cinfo->comp_info[i].v_samp_factor;
+           j++, mem += alignedPlaneWidth[i]) {
+        mcuRowsTmp[i][j] = mem;
+        if (i > 0) {
+          memset(mem + mPlaneWidth[i], 128, alignedPlaneWidth[i] - mPlaneWidth[i]);
+        }
+      }
+    } else if (mPlaneHeight[i] % DCTSIZE != 0) {
+      mPlanesMCURow[i] = std::make_unique<uint8_t[]>(alignedPlaneWidth[i]);
+      if (i > 0) {
+        memset(mPlanesMCURow[i].get(), 128, alignedPlaneWidth[i]);
+      }
     }
+    subImage[i] = strides[i] < alignedPlaneWidth[i] ? mcuRowsTmp[i] : mcuRows[i];
   }
 
   while (cinfo->next_scanline < cinfo->image_height) {
-    for (int i = 0; i < kCompressBatchSize; ++i) {
-      size_t scanline = cinfo->next_scanline + i;
-      if (scanline < cinfo->image_height) {
-        y[i] = y_plane + scanline * lumaStride;
-        if (need_luma_padding) {
-          uint8_t* tmp = y_mcu_row_ptr + i * aligned_luma_width;
-          memcpy(tmp, y[i], cinfo->image_width);
-          y[i] = tmp;
+    JDIMENSION mcu_scanline_start[kMaxNumComponents];
+
+    for (int i = 0; i < cinfo->num_components; i++) {
+      mcu_scanline_start[i] =
+          std::ceil(((float)cinfo->next_scanline * cinfo->comp_info[i].v_samp_factor) /
+                    cinfo->max_v_samp_factor);
+
+      for (int j = 0; j < cinfo->comp_info[i].v_samp_factor * DCTSIZE; j++) {
+        JDIMENSION scanline = mcu_scanline_start[i] + j;
+
+        if (scanline < mPlaneHeight[i]) {
+          mcuRows[i][j] = const_cast<uint8_t*>(planes[i] + scanline * strides[i]);
+          if (strides[i] < alignedPlaneWidth[i]) {
+            memcpy(mcuRowsTmp[i][j], mcuRows[i][j], mPlaneWidth[i]);
+          }
+        } else {
+          mcuRows[i][j] = mPlanesMCURow[i].get();
         }
-      } else {
-        y[i] = empty.get();
       }
     }
-    int processed = jpeg_write_raw_data(cinfo, planes, kCompressBatchSize);
-    if (processed != kCompressBatchSize / 2) {
-      ALOGE("Number of processed lines does not equal input lines.");
+    int processed = jpeg_write_raw_data(cinfo, subImage, DCTSIZE * cinfo->max_v_samp_factor);
+    if (processed != DCTSIZE * cinfo->max_v_samp_factor) {
+      ALOGE("number of scan lines processed %d does not equal requested scan lines %d ", processed,
+            DCTSIZE * cinfo->max_v_samp_factor);
       return false;
     }
   }
diff --git a/lib/src/jpegr.cpp b/lib/src/jpegr.cpp
index ef60113..c755ed7 100644
--- a/lib/src/jpegr.cpp
+++ b/lib/src/jpegr.cpp
@@ -225,7 +225,7 @@ status_t JpegR::encodeJPEGR(jr_uncompressed_ptr p010_image_ptr, ultrahdr_transfe
     p010_image.chroma_stride = p010_image.luma_stride;
   }
 
-  const size_t yu420_luma_stride = ALIGNM(p010_image.width, JpegEncoderHelper::kCompressBatchSize);
+  const size_t yu420_luma_stride = ALIGNM(p010_image.width, 16);
   unique_ptr<uint8_t[]> yuv420_image_data =
       make_unique<uint8_t[]>(yu420_luma_stride * p010_image.height * 3 / 2);
   jpegr_uncompressed_struct yuv420_image;
@@ -269,11 +269,15 @@ status_t JpegR::encodeJPEGR(jr_uncompressed_ptr p010_image_ptr, ultrahdr_transfe
 
   // compress 420 image
   JpegEncoderHelper jpeg_enc_obj_yuv420;
-  if (!jpeg_enc_obj_yuv420.compressImage(reinterpret_cast<uint8_t*>(yuv420_image.data),
-                                         reinterpret_cast<uint8_t*>(yuv420_image.chroma_data),
-                                         yuv420_image.width, yuv420_image.height,
-                                         yuv420_image.luma_stride, yuv420_image.chroma_stride,
-                                         quality, icc->getData(), icc->getLength())) {
+  const uint8_t* planes[3]{reinterpret_cast<uint8_t*>(yuv420_image.data),
+                           reinterpret_cast<uint8_t*>(yuv420_image.chroma_data),
+                           reinterpret_cast<uint8_t*>(yuv420_image.chroma_data) +
+                               yuv420_image.chroma_stride * yuv420_image.height / 2};
+  const size_t strides[3]{yuv420_image.luma_stride, yuv420_image.chroma_stride,
+                          yuv420_image.chroma_stride};
+  if (!jpeg_enc_obj_yuv420.compressImage(planes, strides, yuv420_image.width, yuv420_image.height,
+                                         JpegEncoderHelper::YUV420, quality, icc->getData(),
+                                         icc->getLength())) {
     return ERROR_JPEGR_ENCODE_ERROR;
   }
   jpegr_compressed_struct jpeg;
@@ -344,8 +348,7 @@ status_t JpegR::encodeJPEGR(jr_uncompressed_ptr p010_image_ptr,
   unique_ptr<uint8_t[]> yuv_420_bt601_data;
   // Convert to bt601 YUV encoding for JPEG encode
   if (yuv420_image.colorGamut != ULTRAHDR_COLORGAMUT_P3) {
-    const size_t yuv_420_bt601_luma_stride =
-        ALIGNM(yuv420_image.width, JpegEncoderHelper::kCompressBatchSize);
+    const size_t yuv_420_bt601_luma_stride = ALIGNM(yuv420_image.width, 16);
     yuv_420_bt601_data =
         make_unique<uint8_t[]>(yuv_420_bt601_luma_stride * yuv420_image.height * 3 / 2);
     yuv420_bt601_image.data = yuv_420_bt601_data.get();
@@ -405,11 +408,15 @@ status_t JpegR::encodeJPEGR(jr_uncompressed_ptr p010_image_ptr,
 
   // compress 420 image
   JpegEncoderHelper jpeg_enc_obj_yuv420;
-  if (!jpeg_enc_obj_yuv420.compressImage(
-          reinterpret_cast<uint8_t*>(yuv420_bt601_image.data),
-          reinterpret_cast<uint8_t*>(yuv420_bt601_image.chroma_data), yuv420_bt601_image.width,
-          yuv420_bt601_image.height, yuv420_bt601_image.luma_stride,
-          yuv420_bt601_image.chroma_stride, quality, icc->getData(), icc->getLength())) {
+  const uint8_t* planes[3]{reinterpret_cast<uint8_t*>(yuv420_bt601_image.data),
+                           reinterpret_cast<uint8_t*>(yuv420_bt601_image.chroma_data),
+                           reinterpret_cast<uint8_t*>(yuv420_bt601_image.chroma_data) +
+                               yuv420_bt601_image.chroma_stride * yuv420_bt601_image.height / 2};
+  const size_t strides[3]{yuv420_bt601_image.luma_stride, yuv420_bt601_image.chroma_stride,
+                          yuv420_bt601_image.chroma_stride};
+  if (!jpeg_enc_obj_yuv420.compressImage(planes, strides, yuv420_bt601_image.width,
+                                         yuv420_bt601_image.height, JpegEncoderHelper::YUV420,
+                                         quality, icc->getData(), icc->getLength())) {
     return ERROR_JPEGR_ENCODE_ERROR;
   }
 
@@ -801,18 +808,20 @@ status_t JpegR::compressGainMap(jr_uncompressed_ptr gainmap_image_ptr,
     return ERROR_JPEGR_BAD_PTR;
   }
 
+  const uint8_t* planes[]{reinterpret_cast<uint8_t*>(gainmap_image_ptr->data)};
   if (kUseMultiChannelGainMap) {
-    if (!jpeg_enc_obj_ptr->compressImage(reinterpret_cast<uint8_t*>(gainmap_image_ptr->data),
-                                         gainmap_image_ptr->width, gainmap_image_ptr->height,
+    const size_t strides[]{gainmap_image_ptr->width * 3};
+    if (!jpeg_enc_obj_ptr->compressImage(planes, strides, gainmap_image_ptr->width,
+                                         gainmap_image_ptr->height, JpegEncoderHelper::RGB,
                                          kMapCompressQuality, nullptr, 0)) {
       return ERROR_JPEGR_ENCODE_ERROR;
     }
   } else {
+    const size_t strides[]{gainmap_image_ptr->width};
     // Don't need to convert YUV to Bt601 since single channel
-    if (!jpeg_enc_obj_ptr->compressImage(reinterpret_cast<uint8_t*>(gainmap_image_ptr->data),
-                                         nullptr, gainmap_image_ptr->width,
-                                         gainmap_image_ptr->height, gainmap_image_ptr->luma_stride,
-                                         0, kMapCompressQuality, nullptr, 0)) {
+    if (!jpeg_enc_obj_ptr->compressImage(planes, strides, gainmap_image_ptr->width,
+                                         gainmap_image_ptr->height, JpegEncoderHelper::GRAYSCALE,
+                                         kMapCompressQuality, nullptr, 0)) {
       return ERROR_JPEGR_ENCODE_ERROR;
     }
   }