cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

Browse Source 
PiperOrigin-RevId: 571087863
This commit is contained in:
Copybara-Service 2023-10-05 11:53:31 -07:00
commit 24da737272
3 changed files with 194 additions and 89 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mediapipe/tasks/ios/vision/core/utils/BUILD
View File

@ -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",

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

@ -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
* converted to an RGB format. In case of grayscale images, the mono channel is duplicated in the R,
* G, B channels.
* Irrespective of whether the underlying buffer is grayscale, RGBA, BGRA etc., the `MPPImage` is
* converted to an RGBA 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
* 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

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

@ -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
withWidth:(size_t)width
height:(size_t)height
stride:(size_t)stride
pixelBufferFormat:(OSType)pixelBufferFormatType
error:(NSError **)error;
+ (std::unique_ptr<ImageFrame>)imageFrameFromPixelData:(uint8_t *)pixelData
withWidth:(size_t)width
height:(size_t)height
stride:(size_t)stride
pixelBufferFormat:(OSType)pixelBufferFormatType
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
withWidth:(size_t)width
height:(size_t)height
stride:(size_t)stride
pixelBufferFormat:(OSType)pixelBufferFormatType
error:(NSError **)error {
NSInteger destinationChannelCount = 3;
+ (std::unique_ptr<ImageFrame>)imageFrameFromPixelData:(uint8_t *)pixelData
withWidth:(size_t)width
height:(size_t)height
stride:(size_t)stride
pixelBufferFormat:(OSType)pixelBufferFormatType
error:(NSError **)error {
NSInteger destinationChannelCount = 4;
size_t destinationBytesPerRow = width * destinationChannelCount;
uint8_t *destPixelBufferAddress =
(uint8_t *)[MPPCommonUtils mallocWithSize:sizeof(uint8_t) * height * destinationBytesPerRow
error:error];
if (!destPixelBufferAddress) {
return NULL;
}
ImageFormat::Format imageFormat = ImageFormat::SRGBA;
vImage_Buffer srcBuffer = {.data = pixelData,
.height = (vImagePixelCount)height,
.width = (vImagePixelCount)width,
.rowBytes = stride};
vImage_Buffer destBuffer = {.data = destPixelBufferAddress,
.height = (vImagePixelCount)height,
.width = (vImagePixelCount)width,
.rowBytes = destinationBytesPerRow};
vImage_Buffer destBuffer;
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"];
free(destPixelBufferAddress);
return NULL;
return nullptr;
}
}
if (convertError != kvImageNoError) {
[MPPCommonUtils createCustomError:error
withCode:MPPTasksErrorCodeInternalError
description:@"Image format conversion failed."];
free(destPixelBufferAddress);
return NULL;
description:@"Some error occured while preprocessing the input image. "
@"Please verify that the image is not corrupted."];
return nullptr;
}
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;
size_t destinationBytesPerRow = destinationChannelCount * width;
UInt8 *pixelDataToReturn = NULL;
std::unique_ptr<ImageFrame> imageFrame = nullptr;
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
withWidth:width
height:height
stride:bytesPerRow
pixelBufferFormat:kCVPixelFormatType_32RGBA
error:error];
imageFrame = [MPPPixelDataUtils imageFrameFromPixelData:srcData
withWidth:width
height:height
stride:bytesPerRow
pixelBufferFormat:kCVPixelFormatType_32RGBA
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>(
mediapipe::ImageFormat::SRGB, (int)width, (int)height, (int)destinationBytesPerRow,
static_cast<uint8 *>(pixelDataToReturn),
/*deleter=*/free);
CGDataProviderRef provider = CGDataProviderCreateWithData(
destBuffer.data, destBuffer.data,
internalImageFrame->WidthStep() * internalImageFrame->Height(), callback);
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: {