Merge pull request #4840 from priankakariatyml:ios-image-utils-updates

PiperOrigin-RevId: 571087863

mediapipe/tasks/ios/vision/core/utils/BUILD

@@ -12,6 +12,7 @@ objc_library(
     ],
     module_name = "MPPImageUtils",
     deps = [
+        "//mediapipe/framework/formats:image",
         "//mediapipe/framework/formats:image_format_cc_proto",
         "//mediapipe/framework/formats:image_frame",
         "//mediapipe/tasks/ios/common:MPPCommon",

mediapipe/tasks/ios/vision/core/utils/sources/MPPImage+Utils.h

@@ -14,7 +14,9 @@
 
 #import <Foundation/Foundation.h>
 
+#include "mediapipe/framework/formats/image.h"
 #include "mediapipe/framework/formats/image_frame.h"
+
 #import "mediapipe/tasks/ios/vision/core/sources/MPPImage.h"
 
 NS_ASSUME_NONNULL_BEGIN
@@ -25,9 +27,9 @@ NS_ASSUME_NONNULL_BEGIN
 @interface MPPImage (Utils)
 /**
  * Converts the `MPPImage` into a `mediapipe::ImageFrame`.
- * Irrespective of whether the underlying buffer is grayscale, RGB, RGBA, BGRA etc., the MPPImage is
+ * Irrespective of whether the underlying buffer is grayscale, RGBA, BGRA etc., the `MPPImage` is
- * converted to an RGB format. In case of grayscale images, the mono channel is duplicated in the R,
+ * converted to an RGBA format. In case of grayscale images, the mono channel is duplicated in the
- * G, B channels.
+ * R, G, B channels.
  *
  * @param error Pointer to the memory location where errors if any should be saved. If @c NULL, no
  * error will be saved.
@@ -36,6 +38,32 @@ NS_ASSUME_NONNULL_BEGIN
  */
 - (std::unique_ptr<mediapipe::ImageFrame>)imageFrameWithError:(NSError **)error;
 
+/**
+ * Initializes an `MPPImage` object with the pixels from the given `mediapipe::Image` and source
+ * type and orientation from the given source image.
+ *
+ * Only supports initialization from `mediapipe::Image` of format RGBA.
+ * If `shouldCopyPixelData` is set to `YES`, the newly created `MPPImage` stores a reference to a
+ * deep copied pixel data of the given `image`. Since deep copies are expensive, it is recommended
+ * to not set `shouldCopyPixelData` unless the `MPPImage` must outlive the passed in `image`.
+ *
+ * @param image The `mediapipe::Image` whose pixels are used for creating the `MPPImage`.
+ * @param sourceImage The `MPPImage` whose `orientation` and `imageSourceType` are used for creating
+ * the new `MPPImage`.
+ * @param shouldCopyPixelData `BOOL` that determines if the newly created `MPPImage` stores a
+ * reference to a deep copied pixel data of the given `image` or not.
+ *
+ * @param error Pointer to the memory location where errors if any should be saved. If @c NULL, no
+ * error will be saved.
+ *
+ * @return A new `MPImage` instance with the pixels from the given `mediapipe::Image` and meta data
+ * equal to the `sourceImage`. `nil` if there is an error in initializing the `MPPImage`.
+ */
+- (nullable instancetype)initWithCppImage:(mediapipe::Image &)image
+           cloningPropertiesOfSourceImage:(MPPImage *)sourceImage
+                      shouldCopyPixelData:(BOOL)shouldCopyPixelData
+                                    error:(NSError **)error;
+
 @end
 
 NS_ASSUME_NONNULL_END

mediapipe/tasks/ios/vision/core/utils/sources/MPPImage+Utils.mm

@@ -22,20 +22,34 @@
 #import <CoreImage/CoreImage.h>
 #import <CoreVideo/CoreVideo.h>
 
+#include <memory>
+
 #include "mediapipe/framework/formats/image_format.pb.h"
 
 namespace {
+using ::mediapipe::ImageFormat;
 using ::mediapipe::ImageFrame;
+
+vImage_Buffer allocatedVImageBuffer(vImagePixelCount width, vImagePixelCount height,
+                                    size_t rowBytes) {
+  UInt8 *data = new UInt8[height * rowBytes];
+  return {.data = data, .height = height, .width = width, .rowBytes = rowBytes};
 }
 
+static void FreeDataProviderReleaseCallback(void *buffer, const void *data, size_t size) {
+  delete (vImage_Buffer *)buffer;
+}
+
+}  // namespace
+
 @interface MPPPixelDataUtils : NSObject
 
-+ (uint8_t *)rgbPixelDataFromPixelData:(uint8_t *)pixelData
++ (std::unique_ptr<ImageFrame>)imageFrameFromPixelData:(uint8_t *)pixelData
-                             withWidth:(size_t)width
+                                             withWidth:(size_t)width
-                                height:(size_t)height
+                                                height:(size_t)height
-                                stride:(size_t)stride
+                                                stride:(size_t)stride
-                     pixelBufferFormat:(OSType)pixelBufferFormatType
+                                     pixelBufferFormat:(OSType)pixelBufferFormatType
-                                 error:(NSError **)error;
+                                                 error:(NSError **)error;
 
 @end
 
@@ -49,6 +63,9 @@ using ::mediapipe::ImageFrame;
 @interface MPPCGImageUtils : NSObject
 
 + (std::unique_ptr<ImageFrame>)imageFrameFromCGImage:(CGImageRef)cgImage error:(NSError **)error;
++ (CGImageRef)cgImageFromImageFrame:(std::shared_ptr<ImageFrame>)imageFrame
+                shouldCopyPixelData:(BOOL)shouldCopyPixelData
+                              error:(NSError **)error;
 
 @end
 
@@ -60,42 +77,44 @@ using ::mediapipe::ImageFrame;
 
 @implementation MPPPixelDataUtils : NSObject
 
-+ (uint8_t *)rgbPixelDataFromPixelData:(uint8_t *)pixelData
++ (std::unique_ptr<ImageFrame>)imageFrameFromPixelData:(uint8_t *)pixelData
-                             withWidth:(size_t)width
+                                             withWidth:(size_t)width
-                                height:(size_t)height
+                                                height:(size_t)height
-                                stride:(size_t)stride
+                                                stride:(size_t)stride
-                     pixelBufferFormat:(OSType)pixelBufferFormatType
+                                     pixelBufferFormat:(OSType)pixelBufferFormatType
-                                 error:(NSError **)error {
+                                                 error:(NSError **)error {
-  NSInteger destinationChannelCount = 3;
+  NSInteger destinationChannelCount = 4;
   size_t destinationBytesPerRow = width * destinationChannelCount;
 
-  uint8_t *destPixelBufferAddress =
+  ImageFormat::Format imageFormat = ImageFormat::SRGBA;
-      (uint8_t *)[MPPCommonUtils mallocWithSize:sizeof(uint8_t) * height * destinationBytesPerRow
-                                          error:error];
-
-  if (!destPixelBufferAddress) {
-    return NULL;
-  }
 
   vImage_Buffer srcBuffer = {.data = pixelData,
                              .height = (vImagePixelCount)height,
                              .width = (vImagePixelCount)width,
                              .rowBytes = stride};
 
-  vImage_Buffer destBuffer = {.data = destPixelBufferAddress,
+  vImage_Buffer destBuffer;
-                              .height = (vImagePixelCount)height,
-                              .width = (vImagePixelCount)width,
-                              .rowBytes = destinationBytesPerRow};
 
   vImage_Error convertError = kvImageNoError;
 
+  // Convert the raw pixel data to RGBA format and un-premultiply the alpha from the R, G, B values
+  // since MediaPipe C++ APIs only accept un pre-multiplied channels.
   switch (pixelBufferFormatType) {
     case kCVPixelFormatType_32RGBA: {
-      convertError = vImageConvert_RGBA8888toRGB888(&srcBuffer, &destBuffer, kvImageNoFlags);
+      destBuffer = allocatedVImageBuffer((vImagePixelCount)width, (vImagePixelCount)height,
+                                         destinationBytesPerRow);
+      convertError = vImageUnpremultiplyData_RGBA8888(&srcBuffer, &destBuffer, kvImageNoFlags);
       break;
     }
     case kCVPixelFormatType_32BGRA: {
-      convertError = vImageConvert_BGRA8888toRGB888(&srcBuffer, &destBuffer, kvImageNoFlags);
+      const uint8_t permute_map[4] = {2, 1, 0, 3};
+      destBuffer = allocatedVImageBuffer((vImagePixelCount)width, (vImagePixelCount)height,
+                                         destinationBytesPerRow);
+      convertError =
+          vImagePermuteChannels_ARGB8888(&srcBuffer, &destBuffer, permute_map, kvImageNoFlags);
+      if (convertError == kvImageNoError) {
+        convertError = vImageUnpremultiplyData_RGBA8888(&destBuffer, &destBuffer, kvImageNoFlags);
+      }
       break;
     }
     default: {
@@ -103,68 +122,45 @@ using ::mediapipe::ImageFrame;
                                withCode:MPPTasksErrorCodeInvalidArgumentError
                             description:@"Invalid source pixel buffer format. Expecting one of "
                                         @"kCVPixelFormatType_32RGBA, kCVPixelFormatType_32BGRA"];
-
+      return nullptr;
-      free(destPixelBufferAddress);
-      return NULL;
     }
   }
 
   if (convertError != kvImageNoError) {
     [MPPCommonUtils createCustomError:error
                              withCode:MPPTasksErrorCodeInternalError
-                          description:@"Image format conversion failed."];
+                          description:@"Some error occured while preprocessing the input image. "
-
+                                      @"Please verify that the image is not corrupted."];
-    free(destPixelBufferAddress);
+    return nullptr;
-    return NULL;
   }
 
-  return destPixelBufferAddress;
+  // Uses default deleter
+  return std::make_unique<ImageFrame>(imageFormat, width, height, destinationBytesPerRow,
+                                       static_cast<uint8 *>(destBuffer.data));
 }
 
 @end
 
 @implementation MPPCVPixelBufferUtils
 
-+ (std::unique_ptr<ImageFrame>)rgbImageFrameFromCVPixelBuffer:(CVPixelBufferRef)pixelBuffer
-                                                        error:(NSError **)error {
-  CVPixelBufferLockBaseAddress(pixelBuffer, 0);
-
-  size_t width = CVPixelBufferGetWidth(pixelBuffer);
-  size_t height = CVPixelBufferGetHeight(pixelBuffer);
-
-  size_t destinationChannelCount = 3;
-  size_t destinationStride = destinationChannelCount * width;
-
-  uint8_t *rgbPixelData = [MPPPixelDataUtils
-      rgbPixelDataFromPixelData:(uint8_t *)CVPixelBufferGetBaseAddress(pixelBuffer)
-                      withWidth:CVPixelBufferGetWidth(pixelBuffer)
-                         height:CVPixelBufferGetHeight(pixelBuffer)
-                         stride:CVPixelBufferGetBytesPerRow(pixelBuffer)
-              pixelBufferFormat:CVPixelBufferGetPixelFormatType(pixelBuffer)
-                          error:error];
-
-  CVPixelBufferUnlockBaseAddress(pixelBuffer, 0);
-
-  if (!rgbPixelData) {
-    return nullptr;
-  }
-
-  std::unique_ptr<ImageFrame> imageFrame =
-      absl::make_unique<ImageFrame>(::mediapipe::ImageFormat::SRGB, width, height,
-                                    destinationStride, static_cast<uint8 *>(rgbPixelData),
-                                    /*deleter=*/free);
-
-  return imageFrame;
-}
-
 + (std::unique_ptr<ImageFrame>)imageFrameFromCVPixelBuffer:(CVPixelBufferRef)pixelBuffer
                                                      error:(NSError **)error {
   OSType pixelBufferFormat = CVPixelBufferGetPixelFormatType(pixelBuffer);
+  std::unique_ptr<ImageFrame> imageFrame = nullptr;
 
   switch (pixelBufferFormat) {
     case kCVPixelFormatType_32RGBA:
     case kCVPixelFormatType_32BGRA: {
-      return [MPPCVPixelBufferUtils rgbImageFrameFromCVPixelBuffer:pixelBuffer error:error];
+      CVPixelBufferLockBaseAddress(pixelBuffer, 0);
+      imageFrame = [MPPPixelDataUtils
+          imageFrameFromPixelData:(uint8_t *)CVPixelBufferGetBaseAddress(pixelBuffer)
+                        withWidth:CVPixelBufferGetWidth(pixelBuffer)
+                           height:CVPixelBufferGetHeight(pixelBuffer)
+                           stride:CVPixelBufferGetBytesPerRow(pixelBuffer)
+                pixelBufferFormat:pixelBufferFormat
+                            error:error];
+      CVPixelBufferUnlockBaseAddress(pixelBuffer, 0);
+      break;
     }
     default: {
       [MPPCommonUtils createCustomError:error
@@ -175,7 +171,7 @@ using ::mediapipe::ImageFrame;
     }
   }
 
-  return nullptr;
+  return imageFrame;
 }
 
 @end
@@ -190,20 +186,20 @@ using ::mediapipe::ImageFrame;
   NSInteger channelCount = 4;
   size_t bytesPerRow = channelCount * width;
 
-  NSInteger destinationChannelCount = 3;
+  std::unique_ptr<ImageFrame> imageFrame = nullptr;
-  size_t destinationBytesPerRow = destinationChannelCount * width;
-
-  UInt8 *pixelDataToReturn = NULL;
 
   CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
   // iOS infers bytesPerRow if it is set to 0.
   // See https://developer.apple.com/documentation/coregraphics/1455939-cgbitmapcontextcreate
   // But for segmentation test image, this was not the case.
   // Hence setting it to the value of channelCount*width.
-  // kCGImageAlphaNoneSkipLast specifies that Alpha will always be next to B.
+  // kCGImageAlphaPremultipliedLast specifies that Alpha will always be next to B and the R, G, B
+  // values will be pre multiplied with alpha. Images with alpha != 255 are stored with the R, G, B
+  // values premultiplied with alpha by iOS. Hence `kCGImageAlphaPremultipliedLast` ensures all
+  // kinds of images (alpha from 0 to 255) are correctly accounted for by iOS.
   // kCGBitmapByteOrder32Big specifies that R will be stored before B.
   // In combination they signify a pixelFormat of kCVPixelFormatType32RGBA.
-  CGBitmapInfo bitMapinfoFor32RGBA = kCGImageAlphaNoneSkipLast | kCGBitmapByteOrder32Big;
+  CGBitmapInfo bitMapinfoFor32RGBA = kCGImageAlphaPremultipliedLast | kCGBitmapByteOrder32Big;
   CGContextRef context = CGBitmapContextCreate(nil, width, height, bitsPerComponent, bytesPerRow,
                                                colorSpace, bitMapinfoFor32RGBA);
 
@@ -214,12 +210,12 @@ using ::mediapipe::ImageFrame;
     if (srcData) {
       // We have drawn the image as an RGBA image with 8 bitsPerComponent and hence can safely input
       // a pixel format of type kCVPixelFormatType_32RGBA for conversion by vImage.
-      pixelDataToReturn = [MPPPixelDataUtils rgbPixelDataFromPixelData:srcData
+      imageFrame = [MPPPixelDataUtils imageFrameFromPixelData:srcData
-                                                             withWidth:width
+                                                    withWidth:width
-                                                                height:height
+                                                       height:height
-                                                                stride:bytesPerRow
+                                                       stride:bytesPerRow
-                                                     pixelBufferFormat:kCVPixelFormatType_32RGBA
+                                            pixelBufferFormat:kCVPixelFormatType_32RGBA
-                                                                 error:error];
+                                                        error:error];
     }
 
     CGContextRelease(context);
@@ -227,16 +223,76 @@ using ::mediapipe::ImageFrame;
 
   CGColorSpaceRelease(colorSpace);
 
-  if (!pixelDataToReturn) {
+  return imageFrame;
+}
+
++ (CGImageRef)cgImageFromImageFrame:(std::shared_ptr<ImageFrame>)imageFrame
+                shouldCopyPixelData:(BOOL)shouldCopyPixelData
+                              error:(NSError **)error {
+  CGBitmapInfo bitmapInfo = kCGImageAlphaNoneSkipLast | kCGBitmapByteOrderDefault;
+
+  ImageFrame *internalImageFrame = imageFrame.get();
+  size_t channelCount = 4;
+
+  switch (internalImageFrame->Format()) {
+    case ImageFormat::SRGBA: {
+      bitmapInfo = kCGImageAlphaPremultipliedLast | kCGBitmapByteOrder32Big;
+      break;
+    }
+    default:
+      [MPPCommonUtils createCustomError:error
+                               withCode:MPPTasksErrorCodeInternalError
+                            description:@"An internal error occured."];
+      return nullptr;
+  }
+
+  size_t bitsPerComponent = 8;
+
+  vImage_Buffer sourceBuffer = {
+      .data = (void *)internalImageFrame->MutablePixelData(),
+      .height = static_cast<vImagePixelCount>(internalImageFrame->Height()),
+      .width = static_cast<vImagePixelCount>(internalImageFrame->Width()),
+      .rowBytes = static_cast<size_t>(internalImageFrame->WidthStep())};
+
+  vImage_Buffer destBuffer;
+
+  CGDataProviderReleaseDataCallback callback = nullptr;
+
+  if (shouldCopyPixelData) {
+    destBuffer = allocatedVImageBuffer(static_cast<vImagePixelCount>(internalImageFrame->Width()),
+                                       static_cast<vImagePixelCount>(internalImageFrame->Height()),
+                                       static_cast<size_t>(internalImageFrame->WidthStep()));
+    callback = FreeDataProviderReleaseCallback;
+  } else {
+    destBuffer = sourceBuffer;
+  }
+
+  // Pre-multiply the raw pixels from a `mediapipe::Image` before creating a `CGImage` to ensure
+  // that pixels are displayed correctly irrespective of their alpha values.
+  vImage_Error premultiplyError =
+      vImagePremultiplyData_RGBA8888(&sourceBuffer, &destBuffer, kvImageNoFlags);
+
+  if (premultiplyError != kvImageNoError) {
+    [MPPCommonUtils createCustomError:error
+                             withCode:MPPTasksErrorCodeInternalError
+                          description:@"An internal error occured."];
+
     return nullptr;
   }
 
-  std::unique_ptr<ImageFrame> imageFrame = absl::make_unique<ImageFrame>(
+  CGDataProviderRef provider = CGDataProviderCreateWithData(
-      mediapipe::ImageFormat::SRGB, (int)width, (int)height, (int)destinationBytesPerRow,
+      destBuffer.data, destBuffer.data,
-      static_cast<uint8 *>(pixelDataToReturn),
+      internalImageFrame->WidthStep() * internalImageFrame->Height(), callback);
-      /*deleter=*/free);
+  CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
+  CGImageRef cgImageRef =
+      CGImageCreate(internalImageFrame->Width(), internalImageFrame->Height(), bitsPerComponent,
+                    bitsPerComponent * channelCount, internalImageFrame->WidthStep(), colorSpace,
+                    bitmapInfo, provider, nullptr, YES, kCGRenderingIntentDefault);
 
-  return imageFrame;
+  CGDataProviderRelease(provider);
+  CGColorSpaceRelease(colorSpace);
+
+  return cgImageRef;
 }
 
 @end
@@ -277,6 +333,26 @@ using ::mediapipe::ImageFrame;
 
 @implementation MPPImage (Utils)
 
+- (nullable instancetype)initWithCppImage:(mediapipe::Image &)image
+           cloningPropertiesOfSourceImage:(MPPImage *)sourceImage
+                      shouldCopyPixelData:(BOOL)shouldCopyPixelData
+                                    error:(NSError **)error {
+  switch (sourceImage.imageSourceType) {
+    case MPPImageSourceTypeImage: {
+      CGImageRef cgImageRef = [MPPCGImageUtils cgImageFromImageFrame:image.GetImageFrameSharedPtr()
+                                                 shouldCopyPixelData:shouldCopyPixelData
+                                                               error:error];
+      UIImage *image = [UIImage imageWithCGImage:cgImageRef];
+      CGImageRelease(cgImageRef);
+
+      return [self initWithUIImage:image orientation:sourceImage.orientation error:nil];
+    }
+    default:
+      // TODO Implement Other Source Types.
+      return nil;
+  }
+}
+
 - (std::unique_ptr<ImageFrame>)imageFrameWithError:(NSError **)error {
   switch (self.imageSourceType) {
     case MPPImageSourceTypeSampleBuffer: {