From de9acdfa6847443b1f17bf99fe37cfb52af2ddd7 Mon Sep 17 00:00:00 2001 From: Prianka Liz Kariat Date: Tue, 13 Jun 2023 22:17:41 +0530 Subject: [PATCH] Added iOS segmentation mask --- mediapipe/tasks/ios/vision/core/BUILD | 16 ++ .../tasks/ios/vision/core/sources/MPPMask.h | 107 ++++++++++++ .../tasks/ios/vision/core/sources/MPPMask.mm | 157 ++++++++++++++++++ 3 files changed, 280 insertions(+) create mode 100644 mediapipe/tasks/ios/vision/core/sources/MPPMask.h create mode 100644 mediapipe/tasks/ios/vision/core/sources/MPPMask.mm diff --git a/mediapipe/tasks/ios/vision/core/BUILD b/mediapipe/tasks/ios/vision/core/BUILD index a97410e1a..7efa1e7e8 100644 --- a/mediapipe/tasks/ios/vision/core/BUILD +++ b/mediapipe/tasks/ios/vision/core/BUILD @@ -64,3 +64,19 @@ objc_library( "@com_google_absl//absl/status:statusor", ], ) + +objc_library( + name = "MPPMask", + srcs = ["sources/MPPMask.mm"], + hdrs = ["sources/MPPMask.h"], + copts = [ + "-ObjC++", + "-std=c++17", + ], + deps = [ + "//mediapipe/tasks/ios/common:MPPCommon", + "//mediapipe/tasks/ios/common/utils:MPPCommonUtils", + "//mediapipe/tasks/ios/core:MPPTaskRunner", + "//third_party/apple_frameworks:CoreVideo", + ], +) diff --git a/mediapipe/tasks/ios/vision/core/sources/MPPMask.h b/mediapipe/tasks/ios/vision/core/sources/MPPMask.h new file mode 100644 index 000000000..37f253c63 --- /dev/null +++ b/mediapipe/tasks/ios/vision/core/sources/MPPMask.h @@ -0,0 +1,107 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#import +#import + +NS_ASSUME_NONNULL_BEGIN + +/** The underlying type of the segmentation mask. */ +typedef NS_ENUM(NSUInteger, MPPMaskDataType) { + + /** Represents the native `UInt8 *` type. */ + MPPMaskDataTypeUInt8, + + /** Represents the native `float *` type. */ + MPPMaskDataTypeFloat32, + +} NS_SWIFT_NAME(MaskDataType); + +/** + * The wrapper class for MediaPipe segmentation masks. + * + * Masks are stored as `UInt8 *` or `float *` objects. + * Every mask is has an underlying type which can be accessed using `dataType`. You can access the + * mask as any other type using the appropriate properties. For eg:, if the underlying type is + * `MPPMaskDataTypeUInt8`, in addition to accessing the mask using `uint8Array`, you can access + * 'floatArray` to get the float 32 data. The first time you access the data as a type different + * from the underlying type, an expensive type conversion is performed. Subsequent accesses return a + * pointer to the memory location fo the same type converted array. As type conversions can be + * expensive, it is recommended to limit the accesses to data of types different from the underlying + * type. + * + * Masks that are returned from a MediaPipe Tasks are owned by by the underlying C++ Task. If you + * need to extend the lifetime of these objects, you can invoke the `[MPPMask copy:]` method. + */ +NS_SWIFT_NAME(Mask) +@interface MPPMask : NSObject + +/** The width of the mask. */ +@property(nonatomic, readonly) CGFloat width; + +/** The height of the mask. */ +@property(nonatomic, readonly) CGFloat height; + +/** The data type of the mask. */ +@property(nonatomic, readonly) MPPMaskDataType dataType; + +/** + * The pointer to the memory location where the underlying mask as a single channel `UInt8` array is + * stored. + */ +@property(nonatomic, readonly, assign) const UInt8 *uint8Data; + +/** + * The pointer to the memory location where the underlying mask as a single channel float 32 array + * is stored. + */ +@property(nonatomic, readonly, assign) const float *float32Data; + +/** + * Initializes an `MPPMask` object of tyep `MPPMaskDataTypeUInt8` with the given `UInt8*` data, + * width and height. + * + * @param uint8Data A pointer to the memory location of the `UInt8` data array. + * @param width The width of the mask. + * @param height The height of the mask. + * + * @return A new `MPPMask` instance with the given `UInt8*` data, width and height. + */ +- (nullable instancetype)initWithUInt8Data:(const UInt8 *)uint8Data + width:(NSInteger)width + height:(NSInteger)height NS_DESIGNATED_INITIALIZER; + +/** + * Initializes an `MPPMask` object of tyep `MPPMaskDataTypeFloat32` with the given `float*` data, + * width and height. + * + * @param uint8Data A pointer to the memory location of the `float` data array. + * @param width The width of the mask. + * @param height The height of the mask. + * + * @return A new `MPPMask` instance with the given `float*` data, width and height. + */ +- (nullable instancetype)initWithFloat32Data:(const float *)float32Data + width:(NSInteger)width + height:(NSInteger)height + error:(NSError **)error NS_DESIGNATED_INITIALIZER; + +/** Unavailable. */ +- (instancetype)init NS_UNAVAILABLE; + ++ (instancetype)new NS_UNAVAILABLE; + +@end + +NS_ASSUME_NONNULL_END diff --git a/mediapipe/tasks/ios/vision/core/sources/MPPMask.mm b/mediapipe/tasks/ios/vision/core/sources/MPPMask.mm new file mode 100644 index 000000000..cc6332676 --- /dev/null +++ b/mediapipe/tasks/ios/vision/core/sources/MPPMask.mm @@ -0,0 +1,157 @@ +// Copyright 2023 The MediaPipe Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#import "mediapipe/tasks/ios/vision/core/sources/MPPMask.h" +#import "mediapipe/tasks/ios/common/sources/MPPCommon.h" +#import "mediapipe/tasks/ios/common/utils/sources/MPPCommonUtils.h" + +namespace { +template +T *allocateDataPtr(std::unique_ptr &data, size_t length) { + data = std::unique_ptr(new T[length]); + return data.get(); +} + +template +void copyData(const T *destination, const T *source, size_t length) { + memcpy((void *)destination, source, length * sizeof(T)); +} +} // namespace + +@interface MPPMask () { + const UInt8 *_uint8Data; + const float *_float32Data; + std::unique_ptr _allocatedUInt8Data; + std::unique_ptr _allocatedFloat32Data; +} +@end + +@implementation MPPMask + +- (nullable instancetype)initWithWidth:(NSInteger)width + height:(NSInteger)height + dataType:(MPPMaskDataType)dataType + error:(NSError **)error { + if (dataType < MPPMaskDataTypeUInt8 || dataType > MPPMaskDataTypeFloat32) { + [MPPCommonUtils createCustomError:error + withCode:MPPTasksErrorCodeInvalidArgumentError + description:@"Invalid value for data type."]; + return nil; + } + + self = [super init]; + if (self) { + _width = width; + _height = height; + _dataType = dataType; + } + return self; +} + +- (nullable instancetype)initWithUInt8Data:(const UInt8 *)uint8Data + width:(NSInteger)width + height:(NSInteger)height { + self = [self initWithWidth:width height:height dataType:MPPMaskDataTypeUInt8 error:nil]; + if (self) { + _uint8Data = uint8Data; + } + return self; +} + +- (nullable instancetype)initWithFloat32Data:(const float *)float32Data + width:(NSInteger)width + height:(NSInteger)height { + self = [self initWithWidth:width height:height dataType:MPPMaskDataTypeFloat32 error:nil]; + if (self) { + _float32Data = float32Data; + } + return self; +} + +- (instancetype)initWithUInt8DataToCopy:(const UInt8 *)uint8DataToCopy + width:(NSInteger)width + height:(NSInteger)height { + self = [self initWithWidth:width height:height dataType:MPPMaskDataTypeUInt8 error:nil]; + if (self) { + _uint8Data = allocateDataPtr(_allocatedUInt8Data, _width * _height); + copyData(_uint8Data, uint8DataToCopy, _width * _height); + } + return self; +} + +- (instancetype)initWithFloat32DataToCopy:(const float *)float32DataToCopy + width:(NSInteger)width + height:(NSInteger)height { + self = [self initWithWidth:width height:height dataType:MPPMaskDataTypeFloat32 error:nil]; + if (self) { + _float32Data = allocateDataPtr(_allocatedFloat32Data, _width * _height); + copyData(_float32Data, float32DataToCopy, _width * _height); + } + return self; +} + +- (const UInt8 *)uint8Data { + switch (_dataType) { + case MPPMaskDataTypeUInt8: { + return _uint8Data; + } + case MPPMaskDataTypeFloat32: { + if (_allocatedUInt8Data) { + return _allocatedUInt8Data.get(); + } + UInt8 *data = allocateDataPtr(_allocatedUInt8Data, _width * _height); + for (int i = 0; i < _width * _height; i++) { + data[i] = _float32Data[i] * 255; + } + return data; + } + default: + return NULL; + } +} + +- (const float *)float32Data { + switch (_dataType) { + case MPPMaskDataTypeUInt8: { + if (_allocatedFloat32Data) { + return _allocatedFloat32Data.get(); + } + float *data = allocateDataPtr(_allocatedFloat32Data, _width * _height); + for (int i = 0; i < _width * _height; i++) { + data[i] = _uint8Data[i] / 255; + } + return data; + } + case MPPMaskDataTypeFloat32: { + return _float32Data; + } + default: + return NULL; + } +} + +- (id)copyWithZone:(NSZone *)zone { + switch (_dataType) { + case MPPMaskDataTypeUInt8: + return [[MPPMask alloc] initWithUInt8DataToCopy:self.uint8Data + width:self.width + height:self.height]; + case MPPMaskDataTypeFloat32: + return [[MPPMask alloc] initWithFloat32DataToCopy:self.float32Data + width:self.width + height:self.height]; + } +} + +@end