From f51dd05e89cecb33478287a49d2adde88ba60ab8 Mon Sep 17 00:00:00 2001
From: Prianka Liz Kariat <prianka.kariat@ymedialabs.com>
Date: Thu, 25 May 2023 19:13:22 +0530
Subject: [PATCH] Added optional removal of Z position in gesture recognizer
 test util

---
 .../MPPGestureRecognizerTests.m                |  6 ++++--
 .../MPPGestureRecognizerResult+ProtoHelpers.h  |  4 +++-
 .../MPPGestureRecognizerResult+ProtoHelpers.mm | 18 +++++++++++-------
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/mediapipe/tasks/ios/test/vision/gesture_recognizer/MPPGestureRecognizerTests.m b/mediapipe/tasks/ios/test/vision/gesture_recognizer/MPPGestureRecognizerTests.m
index d457f8d18..1a48322b4 100644
--- a/mediapipe/tasks/ios/test/vision/gesture_recognizer/MPPGestureRecognizerTests.m
+++ b/mediapipe/tasks/ios/test/vision/gesture_recognizer/MPPGestureRecognizerTests.m
@@ -98,7 +98,8 @@ static const float kLandmarksErrorTolerance = 0.03f;
 
   return [MPPGestureRecognizerResult
       gestureRecognizerResultsFromTextEncodedProtobufFileWithName:filePath
-                                                     gestureLabel:kExpectedThumbUpLabel];
+                                                     gestureLabel:kExpectedThumbUpLabel
+                                            shouldRemoveZPosition:YES];
 }
 
 + (MPPGestureRecognizerResult *)fistGestureRecognizerResultWithLabel:(NSString *)gestureLabel {
@@ -106,7 +107,8 @@ static const float kLandmarksErrorTolerance = 0.03f;
 
   return [MPPGestureRecognizerResult
       gestureRecognizerResultsFromTextEncodedProtobufFileWithName:filePath
-                                                     gestureLabel:gestureLabel];
+                                                     gestureLabel:gestureLabel
+                                            shouldRemoveZPosition:YES];
 }
 
 - (void)assertMultiHandLandmarks:(NSArray<NSArray<MPPNormalizedLandmark *> *> *)multiHandLandmarks
diff --git a/mediapipe/tasks/ios/test/vision/gesture_recognizer/utils/sources/MPPGestureRecognizerResult+ProtoHelpers.h b/mediapipe/tasks/ios/test/vision/gesture_recognizer/utils/sources/MPPGestureRecognizerResult+ProtoHelpers.h
index 9abea43e2..6bb2e5182 100644
--- a/mediapipe/tasks/ios/test/vision/gesture_recognizer/utils/sources/MPPGestureRecognizerResult+ProtoHelpers.h
+++ b/mediapipe/tasks/ios/test/vision/gesture_recognizer/utils/sources/MPPGestureRecognizerResult+ProtoHelpers.h
@@ -20,7 +20,9 @@ NS_ASSUME_NONNULL_BEGIN
 
 + (MPPGestureRecognizerResult *)
     gestureRecognizerResultsFromTextEncodedProtobufFileWithName:(NSString *)fileName
-                                                   gestureLabel:(NSString *)gestureLabel;
+                                                   gestureLabel:(NSString *)gestureLabel
+                                          shouldRemoveZPosition:(BOOL)removeZPosition;
+
 @end
 
 NS_ASSUME_NONNULL_END
diff --git a/mediapipe/tasks/ios/test/vision/gesture_recognizer/utils/sources/MPPGestureRecognizerResult+ProtoHelpers.mm b/mediapipe/tasks/ios/test/vision/gesture_recognizer/utils/sources/MPPGestureRecognizerResult+ProtoHelpers.mm
index 3c40bd880..39002eb7e 100644
--- a/mediapipe/tasks/ios/test/vision/gesture_recognizer/utils/sources/MPPGestureRecognizerResult+ProtoHelpers.mm
+++ b/mediapipe/tasks/ios/test/vision/gesture_recognizer/utils/sources/MPPGestureRecognizerResult+ProtoHelpers.mm
@@ -33,18 +33,22 @@ using ::mediapipe::tasks::ios::test::vision::utils::get_proto_from_pbtxt;
 
 + (MPPGestureRecognizerResult *)
     gestureRecognizerResultsFromTextEncodedProtobufFileWithName:(NSString *)fileName
-                                                   gestureLabel:(NSString *)gestureLabel {
+                                                   gestureLabel:(NSString *)gestureLabel
+                                          shouldRemoveZPosition:(BOOL)removeZPosition {
   LandmarksDetectionResultProto landmarkDetectionResultProto;
 
   if (!get_proto_from_pbtxt(fileName.cppString, landmarkDetectionResultProto).ok()) {
     return nil;
   }
-  // Remove z position of landmarks, because they are not used in correctness
-  // testing. For video or live stream mode, the z positions varies a lot during
-  // tracking from frame to frame.
-  for (int i = 0; i < landmarkDetectionResultProto.landmarks().landmark().size(); i++) {
-    auto &landmark = *landmarkDetectionResultProto.mutable_landmarks()->mutable_landmark(i);
-    landmark.clear_z();
+
+  if (removeZPosition) {
+    // Remove z position of landmarks, because they are not used in correctness
+    // testing. For video or live stream mode, the z positions varies a lot during
+    // tracking from frame to frame.
+    for (int i = 0; i < landmarkDetectionResultProto.landmarks().landmark().size(); i++) {
+      auto &landmark = *landmarkDetectionResultProto.mutable_landmarks()->mutable_landmark(i);
+      landmark.clear_z();
+    }
   }
 
   ClassificationListProto gesturesProto;