Merge pull request #4538 from priankakariatyml:ios-segmentation-mask

PiperOrigin-RevId: 541944396
This commit is contained in:
Copybara-Service 2023-06-20 09:12:53 -07:00
commit bd3a8d885d
3 changed files with 267 additions and 0 deletions

View File

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

View File

@ -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 <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 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 <NSCopying>
/** 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

View File

@ -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<UInt8[]> _uint8DataPtr;
std::unique_ptr<float[]> _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<UInt8[]>(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<float[]>(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<UInt8[]>(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<float[]>(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