Added iOS segmentation mask

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",
+    ],
+)

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 <CoreVideo/CoreVideo.h>
+#import <Foundation/Foundation.h>
+
+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 <NSCopying>
+
+/** 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

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 <typename T>
+T *allocateDataPtr(std::unique_ptr<T[]> &data, size_t length) {
+  data = std::unique_ptr<T[]>(new T[length]);
+  return data.get();
+}
+
+template <typename T>
+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<UInt8[]> _allocatedUInt8Data;
+  std::unique_ptr<float[]> _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