Added null check for segmentation masks in pose landmarker helper initializer

This commit is contained in:
Prianka Liz Kariat 2023-11-07 09:49:42 +05:30
parent 1d0f3734b4
commit 91095c2d6a
2 changed files with 17 additions and 12 deletions

View File

@ -20,6 +20,8 @@
NS_ASSUME_NONNULL_BEGIN NS_ASSUME_NONNULL_BEGIN
static const int kMicroSecondsPerMillisecond = 1000;
@interface MPPPoseLandmarkerResult (Helpers) @interface MPPPoseLandmarkerResult (Helpers)
/** /**
@ -56,7 +58,7 @@ NS_ASSUME_NONNULL_BEGIN
worldLandmarksProto: worldLandmarksProto:
(const std::vector<::mediapipe::LandmarkList> &)worldLandmarksProto (const std::vector<::mediapipe::LandmarkList> &)worldLandmarksProto
segmentationMasks:(const std::vector<mediapipe::Image> *)segmentationMasks segmentationMasks:(const std::vector<mediapipe::Image> *)segmentationMasks
timestampInMilliSeconds:(NSInteger)timestampInMilliseconds; timestampInMilliseconds:(NSInteger)timestampInMilliseconds;
@end @end

View File

@ -21,8 +21,6 @@ using LandmarkListProto = ::mediapipe::LandmarkList;
using NormalizedLandmarkListProto = ::mediapipe::NormalizedLandmarkList; using NormalizedLandmarkListProto = ::mediapipe::NormalizedLandmarkList;
using ::mediapipe::Image; using ::mediapipe::Image;
using ::mediapipe::Packet; using ::mediapipe::Packet;
static const int kMicroSecondsPerMilliSecond = 1000;
} // namespace } // namespace
@implementation MPPPoseLandmarkerResult (Helpers) @implementation MPPPoseLandmarkerResult (Helpers)
@ -41,7 +39,7 @@ static const int kMicroSecondsPerMilliSecond = 1000;
worldLandmarksProto: worldLandmarksProto:
(const std::vector<LandmarkListProto> &)worldLandmarksProto (const std::vector<LandmarkListProto> &)worldLandmarksProto
segmentationMasks:(const std::vector<Image> *)segmentationMasks segmentationMasks:(const std::vector<Image> *)segmentationMasks
timestampInMilliSeconds:(NSInteger)timestampInMilliseconds { timestampInMilliseconds:(NSInteger)timestampInMilliseconds {
NSMutableArray<NSMutableArray<MPPNormalizedLandmark *> *> *multiplePoseLandmarks = NSMutableArray<NSMutableArray<MPPNormalizedLandmark *> *> *multiplePoseLandmarks =
[NSMutableArray arrayWithCapacity:(NSUInteger)landmarksProto.size()]; [NSMutableArray arrayWithCapacity:(NSUInteger)landmarksProto.size()];
@ -69,6 +67,12 @@ static const int kMicroSecondsPerMilliSecond = 1000;
[multiplePoseWorldLandmarks addObject:worldLandmarks]; [multiplePoseWorldLandmarks addObject:worldLandmarks];
} }
if (!segmentationMasks) {
return [[MPPPoseLandmarkerResult alloc] initWithLandmarks:multiplePoseLandmarks
worldLandmarks:multiplePoseWorldLandmarks
segmentationMasks:nil
timestampInMilliseconds:timestampInMilliseconds];
}
NSMutableArray<MPPMask *> *confidenceMasks = NSMutableArray<MPPMask *> *confidenceMasks =
[NSMutableArray arrayWithCapacity:(NSUInteger)segmentationMasks->size()]; [NSMutableArray arrayWithCapacity:(NSUInteger)segmentationMasks->size()];
@ -83,12 +87,11 @@ static const int kMicroSecondsPerMilliSecond = 1000;
shouldCopy:YES]]; shouldCopy:YES]];
} }
MPPPoseLandmarkerResult *poseLandmarkerResult = return [[MPPPoseLandmarkerResult alloc] initWithLandmarks:multiplePoseLandmarks
[[MPPPoseLandmarkerResult alloc] initWithLandmarks:multiplePoseLandmarks worldLandmarks:multiplePoseWorldLandmarks
worldLandmarks:multiplePoseWorldLandmarks segmentationMasks:confidenceMasks
segmentationMasks:confidenceMasks timestampInMilliseconds:timestampInMilliseconds];
timestampInMilliseconds:timestampInMilliseconds]; ;
return poseLandmarkerResult;
} }
+ (MPPPoseLandmarkerResult *) + (MPPPoseLandmarkerResult *)
@ -96,7 +99,7 @@ static const int kMicroSecondsPerMilliSecond = 1000;
worldLandmarksPacket:(const Packet &)worldLandmarksPacket worldLandmarksPacket:(const Packet &)worldLandmarksPacket
segmentationMasksPacket:(const Packet *)segmentationMasksPacket { segmentationMasksPacket:(const Packet *)segmentationMasksPacket {
NSInteger timestampInMilliseconds = NSInteger timestampInMilliseconds =
(NSInteger)(landmarksPacket.Timestamp().Value() / kMicroSecondsPerMilliSecond); (NSInteger)(landmarksPacket.Timestamp().Value() / kMicroSecondsPerMillisecond);
if (landmarksPacket.IsEmpty()) { if (landmarksPacket.IsEmpty()) {
return [MPPPoseLandmarkerResult return [MPPPoseLandmarkerResult
@ -118,7 +121,7 @@ static const int kMicroSecondsPerMilliSecond = 1000;
worldLandmarksProto:worldLandmarksPacket worldLandmarksProto:worldLandmarksPacket
.Get<std::vector<LandmarkListProto>>() .Get<std::vector<LandmarkListProto>>()
segmentationMasks:segmentationMasks segmentationMasks:segmentationMasks
timestampInMilliSeconds:timestampInMilliseconds]; timestampInMilliseconds:timestampInMilliseconds];
} }
@end @end