Updated MPPImageUtils with methods to create image frame

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

@@ -4,23 +4,23 @@ licenses(["notice"])
 
 objc_library(
     name = "MPPImageUtils",
-    srcs = ["sources/MPPImage+Utils.m"],
+    srcs = ["sources/MPPImage+Utils.mm"],
     hdrs = ["sources/MPPImage+Utils.h"],
-    copts = [
-        "-ObjC++",
-        "-std=c++17",
-    ],
     module_name = "MPPImageUtils",
     sdk_frameworks = [
         "Accelerate",
         "CoreGraphics",
         "CoreImage",
         "CoreVideo",
-        "UIKit",
     ],
     deps = [
         "//mediapipe/tasks/ios/common/utils:MPPCommonUtils",
         "//mediapipe/tasks/ios/vision/core:MPPImage",
-        "//third_party/apple_frameworks:UIKit",
+        "//mediapipe/framework/formats:image_format_cc_proto",
+        "//mediapipe/framework/formats:image_frame",
+    ],
+    copts = [
+        "-ObjC++",
+        "-std=c++17",
     ],
 )

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

@@ -15,29 +15,25 @@
 #import <Foundation/Foundation.h>
 
 #import "mediapipe/tasks/ios/vision/core/sources/MPPImage.h"
+#include "mediapipe/framework/formats/image_frame.h"
 
 NS_ASSUME_NONNULL_BEGIN
 
-/**
+/** 
- * Helper utility for performing operations on MPPImage specific to the MediaPipe Vision library.
+ * Helper utility for converting `MPPImage` into a `mediapipe::ImageFrame`.
  */
 @interface MPPImage (Utils)
-
-/** Bitmap size of the image. */
-@property(nonatomic, readonly) CGSize bitmapSize;
-
 /**
- * Returns the underlying uint8 pixel buffer of an `MPPImage`.
+ * Converts the `MPPImage` into a `mediapipe::ImageFrame`.
- * Irrespective of whether the underlying buffer is grayscale, RGB, RGBA, BGRA etc., the pixel
+ * Irrespective of whether the underlying buffer is grayscale, RGB, RGBA, BGRA etc., the MPPImage is converted to an RGB format. In case of grayscale images, the mono channel is duplicated
- * data is converted to an RGB format. In case of grayscale images, the mono channel is duplicated
  * in the R, G, B channels.
  *
- * @param error Pointer to the memory location where errors if any should be saved. If @c NULL, no
+ * @param error Pointer to the memory location where errors if any should be
- * error will be saved.
+ * saved. If @c NULL, no error will be saved.
  *
- * @return The underlying pixel buffer of the `MPPImage` or nil in case of errors.
+ * @return An std::unique_ptr<mediapipe::ImageFrame> or `nullptr` in case of errors.
  */
-- (nullable uint8_t *)rgbPixelDataWithError:(NSError **)error;
+- (std::unique_ptr<mediapipe::ImageFrame>)imageFrameWithError:(NSError **)error;
 
 @end
 

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

@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#import "mediapipe/tasks/ios/vision/core/utils/sources/MPPImage+Utils.h"
+#import "mediapipe/tasks/ios/vision/core/utils/sources/MPPImage+ImageFrameUtils.h"
 
 #import "mediapipe/tasks/ios/common/sources/MPPCommon.h"
 #import "mediapipe/tasks/ios/common/utils/sources/MPPCommonUtils.h"
@@ -22,6 +22,12 @@
 #import <CoreImage/CoreImage.h>
 #import <CoreVideo/CoreVideo.h>
 
+#include "mediapipe/framework/formats/image_format.pb.h"
+
+namespace {
+  using ::mediapipe::ImageFrame;
+}
+
 @interface MPPPixelDataUtils : NSObject
 
 + (uint8_t *)rgbPixelDataFromPixelData:(uint8_t *)pixelData
@@ -35,21 +41,19 @@
 
 @interface MPPCVPixelBufferUtils : NSObject
 
-+ (uint8_t *)pixelDataFromCVPixelBuffer:(CVPixelBufferRef)pixelBuffer error:(NSError **)error;
++ (std::unique_ptr<ImageFrame>)imageFrameFromCVPixelBuffer:(CVPixelBufferRef)pixelBuffer error:(NSError **)error;
 
 @end
 
 @interface MPPCGImageUtils : NSObject
 
-+ (UInt8 *_Nullable)pixelDataFromCGImage:(CGImageRef)cgImage error:(NSError **)error;
++ (std::unique_ptr<ImageFrame>)imageFrameFromCGImage:(CGImageRef)cgImage error:(NSError **)error;
 
 @end
 
-@interface UIImage (RawPixelDataUtils)
+@interface UIImage (ImageFrameUtils)
 
-@property(nonatomic, readonly) CGSize bitmapSize;
+- (std::unique_ptr<ImageFrame>)imageFrameWithError:(NSError **)error;
-
-- (uint8_t *)pixelDataWithError:(NSError **)error;
 
 @end
 
@@ -120,9 +124,15 @@
 
 @implementation MPPCVPixelBufferUtils
 
-+ (uint8_t *)rgbPixelDataFromCVPixelBuffer:(CVPixelBufferRef)pixelBuffer error:(NSError **)error {
++ (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 stride = CVPixelBufferGetBytesPerRow(pixelBuffer);
+
+
+
   uint8_t *rgbPixelData = [MPPPixelDataUtils
       rgbPixelDataFromPixelData:(uint8_t *)CVPixelBufferGetBaseAddress(pixelBuffer)
                       withWidth:CVPixelBufferGetWidth(pixelBuffer)
@@ -133,10 +143,19 @@
 
   CVPixelBufferUnlockBaseAddress(pixelBuffer, 0);
 
-  return rgbPixelData;
+  if (!rgbPixelData) {
+    return nullptr;
+  }
+  
+  std::unique_ptr<ImageFrame> imageFrame = absl::make_unique<ImageFrame>(
+      ::mediapipe::ImageFormat::SRGB, /*width=*/width, /*height=*/height, stride,
+      static_cast<uint8*>(rgbPixelData),
+      /*deleter=*/free);
+
+  return imageFrame;
 }
 
-+ (nullable uint8_t *)pixelDataFromCVPixelBuffer:(CVPixelBufferRef)pixelBuffer
++ (std::unique_ptr<ImageFrame>)imageFrameFromCVPixelBuffer:(CVPixelBufferRef)pixelBuffer
                                            error:(NSError **)error {
   uint8_t *pixelData = NULL;
 
@@ -144,8 +163,7 @@
 
   switch (pixelBufferFormat) {
     case kCVPixelFormatType_32BGRA: {
-      pixelData = [MPPCVPixelBufferUtils rgbPixelDataFromCVPixelBuffer:pixelBuffer error:error];
+      return [MPPCVPixelBufferUtils rgbImageFrameFromCVPixelBuffer:pixelBuffer error:error];
-      break;
     }
     default: {
       [MPPCommonUtils createCustomError:error
@@ -155,20 +173,20 @@
     }
   }
 
-  return pixelData;
+  return nullptr;
 }
 
 @end
 
 @implementation MPPCGImageUtils
 
-+ (UInt8 *_Nullable)pixelDataFromCGImage:(CGImageRef)cgImage error:(NSError **)error {
++ (std::unique_ptr<ImageFrame>)imageFrameFromCGImage:(CGImageRef)cgImage error:(NSError **)error {
   size_t width = CGImageGetWidth(cgImage);
   size_t height = CGImageGetHeight(cgImage);
 
   NSInteger bitsPerComponent = 8;
   NSInteger channelCount = 4;
-  UInt8 *pixel_data_to_return = NULL;
+  UInt8 *pixelDataToReturn = NULL;
 
   CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();
   size_t bytesPerRow = channelCount * width;
@@ -191,7 +209,7 @@
     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.
-      pixel_data_to_return = [MPPPixelDataUtils rgbPixelDataFromPixelData:srcData
+      pixelDataToReturn = [MPPPixelDataUtils rgbPixelDataFromPixelData:srcData
                                                                 withWidth:width
                                                                    height:height
                                                                    stride:bytesPerRow
@@ -204,38 +222,42 @@
 
   CGColorSpaceRelease(colorSpace);
 
-  return pixel_data_to_return;
+  std::unique_ptr<ImageFrame> imageFrame = absl::make_unique<ImageFrame>(
+      mediapipe::ImageFormat::SRGB, /*width=*/(int)width, /*height=*/(int)height, (int)bytesPerRow,
+      static_cast<uint8*>(pixelDataToReturn),
+      /*deleter=*/free);
+
+  return imageFrame;
 }
 
 @end
 
-@implementation UIImage (RawPixelDataUtils)
+@implementation UIImage (ImageFrameUtils)
 
-- (uint8_t *)pixelDataFromCIImageWithError:(NSError **)error {
+- (std::unique_ptr<ImageFrame>)imageFrameFromCIImageWithError:(NSError **)error {
-  uint8_t *pixelData = NULL;
 
   if (self.CIImage.pixelBuffer) {
-    pixelData = [MPPCVPixelBufferUtils pixelDataFromCVPixelBuffer:self.CIImage.pixelBuffer
+    return [MPPCVPixelBufferUtils imageFrameFromCVPixelBuffer:self.CIImage.pixelBuffer
                                                             error:error];
 
   } else if (self.CIImage.CGImage) {
-    pixelData = [MPPCGImageUtils pixelDataFromCGImage:self.CIImage.CGImage error:error];
+    return [MPPCGImageUtils imageFrameFromCGImage:self.CIImage.CGImage error:error];
   } else {
     [MPPCommonUtils createCustomError:error
                              withCode:MPPTasksErrorCodeInvalidArgumentError
                           description:@"CIImage should have CGImage or CVPixelBuffer info."];
   }
 
-  return pixelData;
+  return nullptr;
 }
 
-- (uint8_t *)pixelDataWithError:(NSError **)error {
+- (std::unique_ptr<ImageFrame>)imageFrameWithError:(NSError **)error {
   uint8_t *pixelData = nil;
 
   if (self.CGImage) {
-    pixelData = [MPPCGImageUtils pixelDataFromCGImage:self.CGImage error:error];
+    return [MPPCGImageUtils imageFrameFromCGImage:self.CGImage error:error];
   } else if (self.CIImage) {
-    pixelData = [self pixelDataFromCIImageWithError:error];
+    return [self imageFrameFromCIImageWithError:error];
   } else {
     [MPPCommonUtils createCustomError:error
                              withCode:MPPTasksErrorCodeInvalidArgumentError
@@ -243,46 +265,27 @@
                                        " CIImage or CGImage."];
   }
 
-  return pixelData;
+  return nullptr;
 }
 
-- (CGSize)bitmapSize {
-  CGFloat width = 0;
-  CGFloat height = 0;
-
-  if (self.CGImage) {
-    width = CGImageGetWidth(self.CGImage);
-    height = CGImageGetHeight(self.CGImage);
-  } else if (self.CIImage.pixelBuffer) {
-    width = CVPixelBufferGetWidth(self.CIImage.pixelBuffer);
-    height = CVPixelBufferGetHeight(self.CIImage.pixelBuffer);
-  } else if (self.CIImage.CGImage) {
-    width = CGImageGetWidth(self.CIImage.CGImage);
-    height = CGImageGetHeight(self.CIImage.CGImage);
-  }
-  return CGSizeMake(width, height);
-}
 @end
 
 @implementation MPPImage (Utils)
 
-- (nullable uint8_t *)rgbPixelDataWithError:(NSError **)error {
+- (std::unique_ptr<ImageFrame>)imageFrameWithError:(NSError **)error {
   uint8_t *pixelData = NULL;
 
   switch (self.imageSourceType) {
     case MPPImageSourceTypeSampleBuffer: {
       CVPixelBufferRef sampleImagePixelBuffer = CMSampleBufferGetImageBuffer(self.sampleBuffer);
-      pixelData = [MPPCVPixelBufferUtils pixelDataFromCVPixelBuffer:sampleImagePixelBuffer
+      return [MPPCVPixelBufferUtils imageFrameFromCVPixelBuffer:sampleImagePixelBuffer
                                                               error:error];
-      break;
     }
     case MPPImageSourceTypePixelBuffer: {
-      pixelData = [MPPCVPixelBufferUtils pixelDataFromCVPixelBuffer:self.pixelBuffer error:error];
+      return [MPPCVPixelBufferUtils imageFrameFromCVPixelBuffer:self.pixelBuffer error:error];
-      break;
     }
     case MPPImageSourceTypeImage: {
-      pixelData = [self.image pixelDataWithError:error];
+      return [self.image imageFrameWithError:error];
-      break;
     }
     default:
       [MPPCommonUtils createCustomError:error
@@ -290,35 +293,7 @@
                             description:@"Invalid source type for MPPImage."];
   }
 
-  return pixelData;
+  return nullptr;
-}
-
-- (CGSize)bitmapSize {
-  CGFloat width = 0;
-  CGFloat height = 0;
-
-  switch (self.imageSourceType) {
-    case MPPImageSourceTypeSampleBuffer: {
-      CVPixelBufferRef pixelBuffer = CMSampleBufferGetImageBuffer(self.sampleBuffer);
-      width = CVPixelBufferGetWidth(pixelBuffer);
-      height = CVPixelBufferGetHeight(pixelBuffer);
-      break;
-    }
-    case MPPImageSourceTypePixelBuffer: {
-      width = CVPixelBufferGetWidth(self.pixelBuffer);
-      height = CVPixelBufferGetHeight(self.pixelBuffer);
-      break;
-    }
-    case MPPImageSourceTypeImage: {
-      width = self.image.bitmapSize.width;
-      height = self.image.bitmapSize.height;
-      break;
-    }
-    default:
-      break;
-  }
-
-  return CGSizeMake(width, height);
 }
 
 @end