Added iOS segmentation mask
This commit is contained in:
parent
b19b80e10f
commit
de9acdfa68
|
@ -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",
|
||||
],
|
||||
)
|
||||
|
|
107
mediapipe/tasks/ios/vision/core/sources/MPPMask.h
Normal file
107
mediapipe/tasks/ios/vision/core/sources/MPPMask.h
Normal 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
|
157
mediapipe/tasks/ios/vision/core/sources/MPPMask.mm
Normal file
157
mediapipe/tasks/ios/vision/core/sources/MPPMask.mm
Normal 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
|
Loading…
Reference in New Issue
Block a user