Skip to content

Commit

Permalink
Merge pull request #4892 from priankakariatyml:ios-mpp-image-utils-pi…
Browse files Browse the repository at this point in the history 
…xel-buffer-format-fix

PiperOrigin-RevId: 590770032
  • Loading branch information
@copybara-github
copybara-github committed Dec 14, 2023
2 parents 4e89de6 + c1cab6e commit 04bcb0d
Show file tree
Hide file tree
Showing 2 changed files with 203 additions and 75 deletions.
230 changes: 185 additions & 45 deletions mediapipe/tasks/ios/vision/core/utils/sources/MPPImage+Utils.mm
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,20 @@
using ::mediapipe::ImageFormat;
using ::mediapipe::ImageFrame;

vImage_Buffer CreateEmptyVImageBufferFromImageFrame(ImageFrame &imageFrame, bool shouldAllocate) {
UInt8 *data = shouldAllocate ? new UInt8[imageFrame.Height() * imageFrame.WidthStep()] : nullptr;
return {.data = data,
.height = static_cast<vImagePixelCount>(imageFrame.Height()),
.width = static_cast<vImagePixelCount>(imageFrame.Width()),
.rowBytes = static_cast<size_t>(imageFrame.WidthStep())};
}

vImage_Buffer CreateVImageBufferFromImageFrame(ImageFrame &imageFrame) {
vImage_Buffer imageBuffer = CreateEmptyVImageBufferFromImageFrame(imageFrame, false);
imageBuffer.data = imageFrame.MutablePixelData();
return imageBuffer;
}

vImage_Buffer allocatedVImageBuffer(vImagePixelCount width, vImagePixelCount height,
size_t rowBytes) {
UInt8 *data = new UInt8[height * rowBytes];
Expand All @@ -40,6 +54,8 @@ static void FreeDataProviderReleaseCallback(void *buffer, const void *data, size
delete[] (vImage_Buffer *)buffer;
}

static void FreeRefConReleaseCallback(void *refCon, const void *baseAddress) { free(refCon); }

} // namespace

@interface MPPPixelDataUtils : NSObject
Expand All @@ -51,13 +67,35 @@ @interface MPPPixelDataUtils : NSObject
pixelBufferFormat:(OSType)pixelBufferFormatType
error:(NSError **)error;

+ (UInt8 *)pixelDataFromImageFrame:(ImageFrame &)imageFrame
shouldCopy:(BOOL)shouldCopy
error:(NSError **)error;

@end

@interface MPPCVPixelBufferUtils : NSObject

+ (std::unique_ptr<ImageFrame>)imageFrameFromCVPixelBuffer:(CVPixelBufferRef)pixelBuffer
error:(NSError **)error;

// This method is used to create CVPixelBuffer from output images of tasks like `FaceStylizer` only
// when the input `MPImage` source type is `pixelBuffer`.
// Always copies the pixel data of the image frame to the created `CVPixelBuffer`.
//
// The only possible 32 RGBA pixel format of input `CVPixelBuffer` is `kCVPixelFormatType_32BGRA`.
// But Mediapipe does not support inference on images of format `BGRA`. Hence the channels of the
// underlying pixel data of `CVPixelBuffer` are permuted to the supported RGBA format before passing
// them to the task for inference. The pixel format of the output images of any MediaPipe task will
// be the same as the pixel format of the input image. (RGBA in this case).
//
// Since creation of `CVPixelBuffer` from the output image pixels with a format of
// `kCVPixelFormatType_32RGBA` is not possible, the channels of the output C++ image `RGBA` have to
// be permuted to the format `BGRA`. When the pixels are copied to create `CVPixelBuffer` this does
// not pose a challenge.
//
// TODO: Investigate if permuting channels of output `mediapipe::Image` in place is possible for
// creating `CVPixelBuffer`s without copying the underlying pixels.
+ (CVPixelBufferRef)cvPixelBufferFromImageFrame:(ImageFrame &)imageFrame error:(NSError **)error;
@end

@interface MPPCGImageUtils : NSObject
Expand Down Expand Up @@ -99,6 +137,9 @@ @implementation MPPPixelDataUtils : NSObject

// 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.
//
// This method is commonly used for `MPImage`s of all source types. Hence supporting BGRA and RGBA
// formats. Only `pixelBuffer` source type is restricted to `BGRA` format.
switch (pixelBufferFormatType) {
case kCVPixelFormatType_32RGBA: {
destBuffer = allocatedVImageBuffer((vImagePixelCount)width, (vImagePixelCount)height,
Expand All @@ -107,6 +148,8 @@ @implementation MPPPixelDataUtils : NSObject
break;
}
case kCVPixelFormatType_32BGRA: {
// Permute channels to `RGBA` since MediaPipe tasks don't support inference on images of
// format `BGRA`.
const uint8_t permute_map[4] = {2, 1, 0, 3};
destBuffer = allocatedVImageBuffer((vImagePixelCount)width, (vImagePixelCount)height,
destinationBytesPerRow);
Expand All @@ -120,8 +163,7 @@ @implementation MPPPixelDataUtils : NSObject
default: {
[MPPCommonUtils createCustomError:error
withCode:MPPTasksErrorCodeInvalidArgumentError
description:@"Invalid source pixel buffer format. Expecting one of "
@"kCVPixelFormatType_32RGBA, kCVPixelFormatType_32BGRA"];
description:@"Some internal error occured."];
return nullptr;
}
}
Expand All @@ -136,7 +178,47 @@ @implementation MPPPixelDataUtils : NSObject

// Uses default deleter
return std::make_unique<ImageFrame>(imageFormat, width, height, destinationBytesPerRow,
static_cast<uint8 *>(destBuffer.data));
static_cast<uint8 *>(destBuffer.data));
}

+ (UInt8 *)pixelDataFromImageFrame:(ImageFrame &)imageFrame
shouldCopy:(BOOL)shouldCopy
error:(NSError **)error {
vImage_Buffer sourceBuffer = CreateVImageBufferFromImageFrame(imageFrame);

// 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;
vImage_Buffer destinationBuffer;

switch (imageFrame.Format()) {
case ImageFormat::SRGBA: {
destinationBuffer =
shouldCopy ? CreateEmptyVImageBufferFromImageFrame(imageFrame, true) : sourceBuffer;
premultiplyError =
vImagePremultiplyData_RGBA8888(&sourceBuffer, &destinationBuffer, kvImageNoFlags);
break;
}
default: {
[MPPCommonUtils createCustomError:error
withCode:MPPTasksErrorCodeInternalError
description:@"An error occured while processing the output image "
@"pixels of the vision task."];
return nullptr;
}
}

if (premultiplyError != kvImageNoError) {
[MPPCommonUtils
createCustomError:error
withCode:MPPTasksErrorCodeInternalError
description:
@"An error occured while processing the output image pixels of the vision task."];

return nullptr;
}

return (UInt8 *)destinationBuffer.data;
}

@end
Expand All @@ -149,7 +231,8 @@ @implementation MPPCVPixelBufferUtils
std::unique_ptr<ImageFrame> imageFrame = nullptr;

switch (pixelBufferFormat) {
case kCVPixelFormatType_32RGBA:
// Core Video only supports pixel data of order BGRA for 32 bit RGBA images.
// Thus other formats like `kCVPixelFormatType_32BGRA` don't need to be accounted for.
case kCVPixelFormatType_32BGRA: {
CVPixelBufferLockBaseAddress(pixelBuffer, 0);
imageFrame = [MPPPixelDataUtils
Expand All @@ -165,15 +248,58 @@ @implementation MPPCVPixelBufferUtils
default: {
[MPPCommonUtils createCustomError:error
withCode:MPPTasksErrorCodeInvalidArgumentError
description:@"Unsupported pixel format for CVPixelBuffer. Supported "
@"pixel format types are kCVPixelFormatType_32BGRA and "
@"kCVPixelFormatType_32RGBA"];
description:@"Unsupported pixel format for CVPixelBuffer. Expected "
@"kCVPixelFormatType_32BGRA"];
}
}

return imageFrame;
}

+ (CVPixelBufferRef)cvPixelBufferFromImageFrame:(ImageFrame &)imageFrame error:(NSError **)error {
if (imageFrame.Format() != ImageFormat::SRGBA) {
[MPPCommonUtils createCustomError:error
withCode:MPPTasksErrorCodeInternalError
description:@"An error occured while creating a CVPixelBuffer from the "
@"output image of the vision task."];
return nullptr;
}

UInt8 *pixelData = [MPPPixelDataUtils pixelDataFromImageFrame:imageFrame
shouldCopy:YES
error:error];

if (!pixelData) {
return nullptr;
}

const uint8_t permute_map[4] = {2, 1, 0, 3};
vImage_Buffer sourceBuffer = CreateEmptyVImageBufferFromImageFrame(imageFrame, NO);
sourceBuffer.data = pixelData;

if (vImagePermuteChannels_ARGB8888(&sourceBuffer, &sourceBuffer, permute_map, kvImageNoFlags) ==
kvImageNoError) {
CVPixelBufferRef outputBuffer;

OSType pixelBufferFormatType = kCVPixelFormatType_32BGRA;

// Since data is copied, pass in a release callback that will be invoked when the pixel buffer
// is destroyed.
if (CVPixelBufferCreateWithBytes(kCFAllocatorDefault, imageFrame.Width(), imageFrame.Height(),
pixelBufferFormatType, pixelData, imageFrame.WidthStep(),
FreeRefConReleaseCallback, pixelData, nullptr,
&outputBuffer) == kCVReturnSuccess) {
return outputBuffer;
}
}

[MPPCommonUtils createCustomError:error
withCode:MPPTasksErrorCodeInternalError
description:@"An error occured while creating a CVPixelBuffer from the "
@"output image of the vision task."];
return nullptr;
}

@end

@implementation MPPCGImageUtils
Expand Down Expand Up @@ -232,7 +358,14 @@ + (CGImageRef)cgImageFromImageFrame:(std::shared_ptr<ImageFrame>)imageFrame
CGBitmapInfo bitmapInfo = kCGImageAlphaNoneSkipLast | kCGBitmapByteOrderDefault;

ImageFrame *internalImageFrame = imageFrame.get();
size_t channelCount = 4;

UInt8 *pixelData = [MPPPixelDataUtils pixelDataFromImageFrame:*internalImageFrame
shouldCopy:shouldCopyPixelData
error:error];

if (!pixelData) {
return nullptr;
}

switch (internalImageFrame->Format()) {
case ImageFormat::SRGBA: {
Expand All @@ -242,56 +375,41 @@ + (CGImageRef)cgImageFromImageFrame:(std::shared_ptr<ImageFrame>)imageFrame
default:
[MPPCommonUtils createCustomError:error
withCode:MPPTasksErrorCodeInternalError
description:@"An internal error occured."];
description:@"An error occured while creating a CGImage from the "
@"output image of the vision task."];
return nullptr;
}

size_t bitsPerComponent = 8;
CGDataProviderReleaseDataCallback callback = nullptr;

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())};
CGDataProviderRef provider = CGDataProviderCreateWithData(
pixelData, pixelData, internalImageFrame->WidthStep() * internalImageFrame->Height(),
callback);

vImage_Buffer destBuffer;
CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceRGB();

CGDataProviderReleaseDataCallback callback = nullptr;
CGImageRef cgImageRef = 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;
if (provider && colorSpace) {
size_t bitsPerComponent = 8;
size_t channelCount = 4;
cgImageRef =
CGImageCreate(internalImageFrame->Width(), internalImageFrame->Height(), bitsPerComponent,
bitsPerComponent * channelCount, internalImageFrame->WidthStep(), colorSpace,
bitmapInfo, provider, nullptr, YES, kCGRenderingIntentDefault);
}

// 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);
// Can safely pass `NULL` to these functions according to iOS docs.
CGDataProviderRelease(provider);
CGColorSpaceRelease(colorSpace);

if (premultiplyError != kvImageNoError) {
if (!cgImageRef) {
[MPPCommonUtils createCustomError:error
withCode:MPPTasksErrorCodeInternalError
description:@"An internal error occured."];

return nullptr;
description:@"An error occured while converting the output image of the "
@"vision task to a CGImage."];
}

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);

CGDataProviderRelease(provider);
CGColorSpaceRelease(colorSpace);

return cgImageRef;
}

Expand Down Expand Up @@ -347,8 +465,30 @@ - (nullable instancetype)initWithCppImage:(mediapipe::Image &)image

return [self initWithUIImage:image orientation:sourceImage.orientation error:nil];
}
case MPPImageSourceTypePixelBuffer: {
if (!shouldCopyPixelData) {
// TODO: Investigate possibility of permuting channels of `mediapipe::Image` returned by
// vision tasks in place to ensure that we can support creating `CVPixelBuffer`s without
// copying the pixel data.
[MPPCommonUtils
createCustomError:error
withCode:MPPTasksErrorCodeInvalidArgumentError
description:
@"When the source type is pixel buffer, you cannot request uncopied data"];
return nil;
}

CVPixelBufferRef pixelBuffer =
[MPPCVPixelBufferUtils cvPixelBufferFromImageFrame:*(image.GetImageFrameSharedPtr())
error:error];
MPPImage *image = [self initWithPixelBuffer:pixelBuffer
orientation:sourceImage.orientation
error:nil];
CVPixelBufferRelease(pixelBuffer);
return image;
}
default:
// TODO Implement Other Source Types.
// TODO Implement CMSampleBuffer.
return nil;
}
}
Expand Down
Loading

0 comments on commit 04bcb0d

Please sign in to comment.