Added iOS segmentation mask

This commit is contained in:
Prianka Liz Kariat 2023-06-13 22:17:41 +05:30
parent b19b80e10f
commit de9acdfa68
3 changed files with 280 additions and 0 deletions

View File

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

View File

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

View File

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