From 91095c2d6affa536b63fa0df88adb3862d4a54fb Mon Sep 17 00:00:00 2001
From: Prianka Liz Kariat <prianka.kariat@codeandtheory.com>
Date: Tue, 7 Nov 2023 09:49:42 +0530
Subject: [PATCH] Added null check for segmentation masks in pose landmarker
 helper  initializer

---
 .../sources/MPPPoseLandmarkerResult+Helpers.h |  4 ++-
 .../MPPPoseLandmarkerResult+Helpers.mm        | 25 +++++++++++--------
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/mediapipe/tasks/ios/vision/pose_landmarker/utils/sources/MPPPoseLandmarkerResult+Helpers.h b/mediapipe/tasks/ios/vision/pose_landmarker/utils/sources/MPPPoseLandmarkerResult+Helpers.h
index 0b20dc32a..83379f676 100644
--- a/mediapipe/tasks/ios/vision/pose_landmarker/utils/sources/MPPPoseLandmarkerResult+Helpers.h
+++ b/mediapipe/tasks/ios/vision/pose_landmarker/utils/sources/MPPPoseLandmarkerResult+Helpers.h
@@ -20,6 +20,8 @@
 
 NS_ASSUME_NONNULL_BEGIN
 
+static const int kMicroSecondsPerMillisecond = 1000;
+
 @interface MPPPoseLandmarkerResult (Helpers)
 
 /**
@@ -56,7 +58,7 @@ NS_ASSUME_NONNULL_BEGIN
                        worldLandmarksProto:
                            (const std::vector<::mediapipe::LandmarkList> &)worldLandmarksProto
                          segmentationMasks:(const std::vector<mediapipe::Image> *)segmentationMasks
-                   timestampInMilliSeconds:(NSInteger)timestampInMilliseconds;
+                   timestampInMilliseconds:(NSInteger)timestampInMilliseconds;
 
 @end
 
diff --git a/mediapipe/tasks/ios/vision/pose_landmarker/utils/sources/MPPPoseLandmarkerResult+Helpers.mm b/mediapipe/tasks/ios/vision/pose_landmarker/utils/sources/MPPPoseLandmarkerResult+Helpers.mm
index 6cd67ff9c..c83365c86 100644
--- a/mediapipe/tasks/ios/vision/pose_landmarker/utils/sources/MPPPoseLandmarkerResult+Helpers.mm
+++ b/mediapipe/tasks/ios/vision/pose_landmarker/utils/sources/MPPPoseLandmarkerResult+Helpers.mm
@@ -21,8 +21,6 @@ using LandmarkListProto = ::mediapipe::LandmarkList;
 using NormalizedLandmarkListProto = ::mediapipe::NormalizedLandmarkList;
 using ::mediapipe::Image;
 using ::mediapipe::Packet;
-
-static const int kMicroSecondsPerMilliSecond = 1000;
 }  // namespace
 
 @implementation MPPPoseLandmarkerResult (Helpers)
@@ -41,7 +39,7 @@ static const int kMicroSecondsPerMilliSecond = 1000;
                        worldLandmarksProto:
                            (const std::vector<LandmarkListProto> &)worldLandmarksProto
                          segmentationMasks:(const std::vector<Image> *)segmentationMasks
-                   timestampInMilliSeconds:(NSInteger)timestampInMilliseconds {
+                   timestampInMilliseconds:(NSInteger)timestampInMilliseconds {
   NSMutableArray<NSMutableArray<MPPNormalizedLandmark *> *> *multiplePoseLandmarks =
       [NSMutableArray arrayWithCapacity:(NSUInteger)landmarksProto.size()];
 
@@ -69,6 +67,12 @@ static const int kMicroSecondsPerMilliSecond = 1000;
     [multiplePoseWorldLandmarks addObject:worldLandmarks];
   }
 
+  if (!segmentationMasks) {
+    return [[MPPPoseLandmarkerResult alloc] initWithLandmarks:multiplePoseLandmarks
+                                               worldLandmarks:multiplePoseWorldLandmarks
+                                            segmentationMasks:nil
+                                      timestampInMilliseconds:timestampInMilliseconds];
+  }
   NSMutableArray<MPPMask *> *confidenceMasks =
       [NSMutableArray arrayWithCapacity:(NSUInteger)segmentationMasks->size()];
 
@@ -83,12 +87,11 @@ static const int kMicroSecondsPerMilliSecond = 1000;
                                                          shouldCopy:YES]];
   }
 
-  MPPPoseLandmarkerResult *poseLandmarkerResult =
-      [[MPPPoseLandmarkerResult alloc] initWithLandmarks:multiplePoseLandmarks
-                                          worldLandmarks:multiplePoseWorldLandmarks
-                                       segmentationMasks:confidenceMasks
-                                 timestampInMilliseconds:timestampInMilliseconds];
-  return poseLandmarkerResult;
+  return [[MPPPoseLandmarkerResult alloc] initWithLandmarks:multiplePoseLandmarks
+                                             worldLandmarks:multiplePoseWorldLandmarks
+                                          segmentationMasks:confidenceMasks
+                                    timestampInMilliseconds:timestampInMilliseconds];
+  ;
 }
 
 + (MPPPoseLandmarkerResult *)
@@ -96,7 +99,7 @@ static const int kMicroSecondsPerMilliSecond = 1000;
                        worldLandmarksPacket:(const Packet &)worldLandmarksPacket
                     segmentationMasksPacket:(const Packet *)segmentationMasksPacket {
   NSInteger timestampInMilliseconds =
-      (NSInteger)(landmarksPacket.Timestamp().Value() / kMicroSecondsPerMilliSecond);
+      (NSInteger)(landmarksPacket.Timestamp().Value() / kMicroSecondsPerMillisecond);
 
   if (landmarksPacket.IsEmpty()) {
     return [MPPPoseLandmarkerResult
@@ -118,7 +121,7 @@ static const int kMicroSecondsPerMilliSecond = 1000;
                          worldLandmarksProto:worldLandmarksPacket
                                                  .Get<std::vector<LandmarkListProto>>()
                            segmentationMasks:segmentationMasks
-                     timestampInMilliSeconds:timestampInMilliseconds];
+                     timestampInMilliseconds:timestampInMilliseconds];
 }
 
 @end