diff --git a/mediapipe/tasks/ios/vision/core/BUILD b/mediapipe/tasks/ios/vision/core/BUILD index a97410e1..a08b2dd6 100644 --- a/mediapipe/tasks/ios/vision/core/BUILD +++ b/mediapipe/tasks/ios/vision/core/BUILD @@ -64,3 +64,17 @@ 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", + ], +) 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 00000000..2227baaa --- /dev/null +++ b/mediapipe/tasks/ios/vision/core/sources/MPPMask.h @@ -0,0 +1,118 @@ +// 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 + +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 has an underlying type which can be accessed using `dataType`. You can access the + * mask as any other type using the appropriate properties. For example, if the underlying type is + * `MPPMaskDataTypeUInt8`, in addition to accessing the mask using `uint8Data`, you can access + * `float32Data` to get the 32 bit float data (with values ranging from 0.0 to 1.0). 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) NSInteger width; + +/** The height of the mask. */ +@property(nonatomic, readonly) NSInteger 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. Uint8 values use the full value range and range from 0 to 255. + */ +@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. Float values range from 0.0 to 1.0. + */ +@property(nonatomic, readonly, assign) const float *float32Data; + +/** + * Initializes an `MPPMask` object of type `MPPMaskDataTypeUInt8` with the given `UInt8*` data, + * width and height. + * + * If `shouldCopy` is set to `YES`, the newly created `MPPMask` stores a reference to a deep copied + * `uint8Data`. Since deep copies are expensive, it is recommended to not set `shouldCopy` unless + * the `MPPMask` must outlive the passed in `uint8Data`. + * + * @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. + * @param shouldCopy 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 + shouldCopy:(BOOL)shouldCopy NS_DESIGNATED_INITIALIZER; + +/** + * Initializes an `MPPMask` object of type `MPPMaskDataTypeFloat32` with the given `float*` data, + * width and height. + * + * If `shouldCopy` is set to `YES`, the newly created `MPPMask` stores a reference to a deep copied + * `float32Data`. Since deep copies are expensive, it is recommended to not set `shouldCopy` unless + * the `MPPMask` must outlive the passed in `float32Data`. + * + * @param float32Data 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 + shouldCopy:(BOOL)shouldCopy NS_DESIGNATED_INITIALIZER; + +// TODO: Add methods for CVPixelBuffer conversion. + +/** 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 00000000..0d78e11d --- /dev/null +++ b/mediapipe/tasks/ios/vision/core/sources/MPPMask.mm @@ -0,0 +1,135 @@ +// 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" + +@interface MPPMask () { + const UInt8 *_uint8Data; + const float *_float32Data; + std::unique_ptr _uint8DataPtr; + std::unique_ptr _float32DataPtr; +} +@end + +@implementation MPPMask + +- (nullable instancetype)initWithUInt8Data:(const UInt8 *)uint8Data + width:(NSInteger)width + height:(NSInteger)height + shouldCopy:(BOOL)shouldCopy { + + self = [super init]; + if (self) { + _width = width; + _height = height; + _dataType = MPPMaskDataTypeUInt8; + + if (shouldCopy) { + size_t length = _width * _height; + _uint8DataPtr = std::unique_ptr(new UInt8[length]); + _uint8Data = _uint8DataPtr.get(); + memcpy((UInt8 *)_uint8Data, uint8Data, length * sizeof(UInt8)); + } else { + _uint8Data = uint8Data; + } + } + return self; +} + +- (nullable instancetype)initWithFloat32Data:(const float *)float32Data + width:(NSInteger)width + height:(NSInteger)height + shouldCopy:(BOOL)shouldCopy { + self = [super init]; + if (self) { + _width = width; + _height = height; + _dataType = MPPMaskDataTypeFloat32; + + if (shouldCopy) { + size_t length = _width * _height; + _float32DataPtr = std::unique_ptr(new float[length]); + _float32Data = _float32DataPtr.get(); + memcpy((float *)_float32Data, float32Data, length * sizeof(float)); + } else { + _float32Data = float32Data; + } + } + return self; +} + +- (const UInt8 *)uint8Data { + switch (_dataType) { + case MPPMaskDataTypeUInt8: { + return _uint8Data; + } + case MPPMaskDataTypeFloat32: { + if (_uint8DataPtr) { + return _uint8DataPtr.get(); + } + + size_t length = _width * _height; + _uint8DataPtr = std::unique_ptr(new UInt8[length]); + UInt8 *data = _uint8DataPtr.get(); + for (int i = 0; i < length; i++) { + data[i] = _float32Data[i] * 255; + } + return data; + } + default: + return NULL; + } +} + +- (const float *)float32Data { + switch (_dataType) { + case MPPMaskDataTypeUInt8: { + if (_float32DataPtr) { + return _float32DataPtr.get(); + } + + size_t length = _width * _height; + _float32DataPtr = std::unique_ptr(new float[length]); + float *data = _float32DataPtr.get(); + for (int i = 0; i < length; i++) { + data[i] = (float)_uint8Data[i] / 255; + } + return data; + } + case MPPMaskDataTypeFloat32: { + return _float32Data; + } + default: + return NULL; + } +} + +- (id)copyWithZone:(NSZone *)zone { + switch (_dataType) { + case MPPMaskDataTypeUInt8: + return [[MPPMask alloc] initWithUInt8Data:self.uint8Data + width:self.width + height:self.height + shouldCopy:YES]; + case MPPMaskDataTypeFloat32: + return [[MPPMask alloc] initWithFloat32Data:self.float32Data + width:self.width + height:self.height + shouldCopy:YES]; + } +} + +@end