From 81ec5801ea34598d93a3e57bc17c2f7807ddbd43 Mon Sep 17 00:00:00 2001
From: Prianka Liz Kariat <prianka.kariat@codeandtheory.com>
Date: Fri, 15 Sep 2023 14:18:10 +0530
Subject: [PATCH] Added new initializers for iOS MPPImage in test utils

---
 mediapipe/tasks/ios/test/vision/utils/BUILD   |  6 +++-
 .../vision/utils/sources/MPPImage+TestUtils.h | 30 +++++++++++++++++++
 .../vision/utils/sources/MPPImage+TestUtils.m | 24 +++++++++++++++
 3 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/mediapipe/tasks/ios/test/vision/utils/BUILD b/mediapipe/tasks/ios/test/vision/utils/BUILD
index d117ad73d..b1f1640d1 100644
--- a/mediapipe/tasks/ios/test/vision/utils/BUILD
+++ b/mediapipe/tasks/ios/test/vision/utils/BUILD
@@ -7,7 +7,11 @@ objc_library(
     srcs = ["sources/MPPImage+TestUtils.m"],
     hdrs = ["sources/MPPImage+TestUtils.h"],
     module_name = "MPPImageTestUtils",
-    deps = ["//mediapipe/tasks/ios/vision/core:MPPImage"],
+    deps = [
+        "//mediapipe/tasks/ios/vision/core:MPPImage",
+        "//mediapipe/tasks/ios/test/utils:MPPFileInfo"],
+)
+
 )
 
 cc_library(
diff --git a/mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.h b/mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.h
index 8cd1c6a67..585677b3f 100644
--- a/mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.h
+++ b/mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.h
@@ -14,6 +14,7 @@
 
 #import <Foundation/Foundation.h>
 
+#import "mediapipe/tasks/ios/test/utils/sources/MPPFileInfo.h"
 #import "mediapipe/tasks/ios/vision/core/sources/MPPImage.h"
 
 NS_ASSUME_NONNULL_BEGIN
@@ -23,6 +24,34 @@ NS_ASSUME_NONNULL_BEGIN
  */
 @interface MPPImage (TestUtils)
 
+/**
+ * Loads an image from a file in an app bundle into a `MPPImage` object.
+ *
+ * @param fileInfo The file info specifying the name and extension of the image
+ * file in the bundle.
+ *
+ * @return The `MPPImage` object contains the loaded image. This method returns
+ * nil if it cannot load the image.
+ */
++ (MPPImage *)imageWithFileInfo:(MPPFileInfo *)fileInfo;
+
+NS_SWIFT_NAME(image(withFileInfo:));
+
+/**
+ * Loads an image from a file in an app bundle into a `MPPImage` object with the specified
+ * orientation.
+ *
+ * @param fileInfo The file info specifying the name and extension of the image
+ * file in the bundle.
+ *
+ * @return The `MPPImage` object contains the loaded image. This method returns
+ * nil if it cannot load the image.
+ */
++ (MPPImage *)imageWithFileInfo:(MPPFileInfo *)fileInfo
+                    orientation:(UIImageOrientation)orientation
+    NS_SWIFT_NAME(image(withFileInfo:orientation:));
+
+// TODO: Remove after all tests are migrated
 /**
  * Loads an image from a file in an app bundle into a `MPPImage` object.
  *
@@ -39,6 +68,7 @@ NS_ASSUME_NONNULL_BEGIN
                                          ofType:(NSString *)type
     NS_SWIFT_NAME(imageFromBundle(class:filename:type:));
 
+// TODO: Remove after all tests are migrated
 /**
  * Loads an image from a file in an app bundle into a `MPPImage` object with the specified
  * orientation.
diff --git a/mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.m b/mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.m
index 0b0ef9fbf..f922146fc 100644
--- a/mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.m
+++ b/mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.m
@@ -14,6 +14,7 @@
 
 #import "mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.h"
 
+// TODO: Remove this category after all tests are migrated to the new methods.
 @interface UIImage (FileUtils)
 
 + (nullable UIImage *)imageFromBundleWithClass:(Class)classObject
@@ -37,6 +38,28 @@
 
 @implementation MPPImage (TestUtils)
 
++ (MPPImage *)imageWithFileInfo:(MPPFileInfo *)fileInfo {
+  if (!fileInfo.path) return nil;
+
+  UIImage *image = [[UIImage alloc] initWithContentsOfFile:fileInfo.path];
+
+  if (!image) return nil;
+
+  return [[MPPImage alloc] initWithUIImage:image error:nil];
+}
+
++ (MPPImage *)imageWithFileInfo:(MPPFileInfo *)fileInfo
+                    orientation:(UIImageOrientation)orientation {
+  if (!fileInfo.path) return nil;
+
+  UIImage *image = [[UIImage alloc] initWithContentsOfFile:fileInfo.path];
+
+  if (!image) return nil;
+
+  return [[MPPImage alloc] initWithUIImage:image orientation:orientation error:nil];
+}
+
+// TODO: Remove after all tests are migrated
 + (nullable MPPImage *)imageFromBundleWithClass:(Class)classObject
                                        fileName:(NSString *)name
                                          ofType:(NSString *)type {
@@ -45,6 +68,7 @@
   return [[MPPImage alloc] initWithUIImage:image error:nil];
 }
 
+// TODO: Remove after all tests are migrated
 + (nullable MPPImage *)imageFromBundleWithClass:(Class)classObject
                                        fileName:(NSString *)name
                                          ofType:(NSString *)type