Merge pull request #4802 from priankakariatyml:ios-image-segmenter-basic-tests

PiperOrigin-RevId: 565797057

mediapipe/tasks/ios/test/utils/BUILD

@@ -0,0 +1,24 @@
+# Copyright 2023 The MediaPipe Authors. All Rights Reserved.
+#
+# 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.
+
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])
+
+objc_library(
+    name = "MPPFileInfo",
+    srcs = ["sources/MPPFileInfo.m"],
+    hdrs = ["sources/MPPFileInfo.h"],
+    module_name = "MPPFileInfo",
+)

mediapipe/tasks/ios/test/utils/sources/MPPFileInfo.h

@@ -0,0 +1,42 @@
+// 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 <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface MPPFileInfo : NSObject
+
+/** The name of the file. */
+@property(nonatomic, readonly) NSString *name;
+
+/** The type of the file. */
+@property(nonatomic, readonly) NSString *type;
+
+/** The path to file in the app bundle. */
+@property(nonatomic, readonly, nullable) NSString *path;
+
+/**
+ * Initializes an `MPPFileInfo` using the given name and type of file.
+ *
+ * @param name The name of the file.
+ * @param type The type of the file.
+ *
+ * @return The `MPPFileInfo` with the given name and type of file.
+ */
+- (instancetype)initWithName:(NSString *)name type:(NSString *)type;
+
+@end
+
+NS_ASSUME_NONNULL_END

mediapipe/tasks/ios/test/utils/sources/MPPFileInfo.m

@@ -0,0 +1,33 @@
+// Copyright 2023 The TensorFlow 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/test/utils/sources/MPPFileInfo.h"
+
+@implementation MPPFileInfo
+
+- (instancetype)initWithName:(NSString *)name type:(NSString *)type {
+  self = [super init];
+  if (self) {
+    _name = name;
+    _type = type;
+  }
+
+  return self;
+}
+
+- (NSString *)path {
+  return [[NSBundle bundleForClass:self.class] pathForResource:self.name ofType:self.type];
+}
+
+@end

mediapipe/tasks/ios/test/vision/image_segmenter/BUILD

@@ -0,0 +1,76 @@
+# Copyright 2023 The MediaPipe Authors. All Rights Reserved.
+#
+# 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.
+
+load(
+    "//mediapipe/framework/tool:ios.bzl",
+    "MPP_TASK_MINIMUM_OS_VERSION",
+)
+load(
+    "@org_tensorflow//tensorflow/lite:special_rules.bzl",
+    "tflite_ios_lab_runner",
+)
+load("@build_bazel_rules_apple//apple:ios.bzl", "ios_unit_test")
+
+package(default_visibility = ["//visibility:public"])
+
+licenses(["notice"])
+
+# Default tags for filtering iOS targets. Targets are restricted to Apple platforms.
+TFL_DEFAULT_TAGS = [
+    "apple",
+]
+
+# Following sanitizer tests are not supported by iOS test targets.
+TFL_DISABLED_SANITIZER_TAGS = [
+    "noasan",
+    "nomsan",
+    "notsan",
+]
+
+objc_library(
+    name = "MPPImageSegmenterObjcTestLibrary",
+    testonly = 1,
+    srcs = ["MPPImageSegmenterTests.mm"],
+    copts = [
+        "-ObjC++",
+        "-std=c++17",
+        "-x objective-c++",
+    ],
+    data = [
+        "//mediapipe/tasks/testdata/vision:test_images",
+        "//mediapipe/tasks/testdata/vision:test_models",
+        "//mediapipe/tasks/testdata/vision:test_protos",
+    ],
+    deps = [
+        "//mediapipe/tasks/ios/test/vision/utils:MPPImageTestUtils",
+        "//mediapipe/tasks/ios/test/vision/utils:MPPMaskTestUtils",
+        "//mediapipe/tasks/ios/vision/image_segmenter:MPPImageSegmenter",
+        "//mediapipe/tasks/ios/vision/image_segmenter:MPPImageSegmenterResult",
+    ] + select({
+        "//third_party:opencv_ios_sim_arm64_source_build": ["@ios_opencv_source//:opencv_xcframework"],
+        "//third_party:opencv_ios_arm64_source_build": ["@ios_opencv_source//:opencv_xcframework"],
+        "//third_party:opencv_ios_x86_64_source_build": ["@ios_opencv_source//:opencv_xcframework"],
+        "//conditions:default": ["@ios_opencv//:OpencvFramework"],
+    }),
+)
+
+ios_unit_test(
+    name = "MPPImageSegmenterObjcTest",
+    minimum_os_version = MPP_TASK_MINIMUM_OS_VERSION,
+    runner = tflite_ios_lab_runner("IOS_LATEST"),
+    tags = TFL_DEFAULT_TAGS + TFL_DISABLED_SANITIZER_TAGS,
+    deps = [
+        ":MPPImageSegmenterObjcTestLibrary",
+    ],
+)

mediapipe/tasks/ios/test/vision/image_segmenter/MPPImageSegmenterTests.mm

@@ -0,0 +1,242 @@
+// 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 <Foundation/Foundation.h>
+#import <XCTest/XCTest.h>
+
+#import "mediapipe/tasks/ios/test/vision/utils/sources/MPPImage+TestUtils.h"
+#import "mediapipe/tasks/ios/test/vision/utils/sources/MPPMask+TestUtils.h"
+#import "mediapipe/tasks/ios/vision/image_segmenter/sources/MPPImageSegmenter.h"
+#import "mediapipe/tasks/ios/vision/image_segmenter/sources/MPPImageSegmenterResult.h"
+
+#include <iostream>
+#include <vector>
+
+static MPPFileInfo *const kCatImageFileInfo = [[MPPFileInfo alloc] initWithName:@"cat" type:@"jpg"];
+static MPPFileInfo *const kCatGoldenImageFileInfo = [[MPPFileInfo alloc] initWithName:@"cat_mask"
+                                                                                 type:@"jpg"];
+static MPPFileInfo *const kSegmentationImageFileInfo =
+    [[MPPFileInfo alloc] initWithName:@"segmentation_input_rotation0" type:@"jpg"];
+static MPPFileInfo *const kSegmentationGoldenImageFileInfo =
+    [[MPPFileInfo alloc] initWithName:@"segmentation_golden_rotation0" type:@"png"];
+static MPPFileInfo *const kImageSegmenterModel = [[MPPFileInfo alloc] initWithName:@"deeplabv3"
+                                                                              type:@"tflite"];
+static NSString *const kExpectedErrorDomain = @"com.google.mediapipe.tasks";
+constexpr float kSimilarityThreshold = 0.96f;
+constexpr NSInteger kMagnificationFactor = 10;
+
+namespace {
+double sum(const std::vector<float> &mask) {
+  double sum = 0.0;
+  for (const float &maskElement : mask) {
+    sum += maskElement;
+  }
+  return sum;
+}
+
+std::vector<float> multiply(const float *mask1, const float *mask2, size_t size) {
+  std::vector<float> multipliedMask;
+  multipliedMask.reserve(size);
+
+  for (int i = 0; i < size; i++) {
+    multipliedMask.push_back(mask1[i] * mask2[i]);
+  }
+
+  return multipliedMask;
+}
+
+double softIOU(const float *mask1, const float *mask2, size_t size) {
+  std::vector<float> interSectionVector = multiply(mask1, mask2, size);
+  double interSectionSum = sum(interSectionVector);
+
+  std::vector<float> m1m1Vector = multiply(mask1, mask1, size);
+  double m1m1 = sum(m1m1Vector);
+
+  std::vector<float> m2m2Vector = multiply(mask2, mask2, size);
+  double m2m2 = sum(m2m2Vector);
+
+  double unionSum = m1m1 + m2m2 - interSectionSum;
+
+  return unionSum > 0.0 ? interSectionSum / unionSum : 0.0;
+}
+} // namespace
+
+@interface MPPImageSegmenterTests : XCTestCase <MPPImageSegmenterLiveStreamDelegate>
+
+@end
+
+@implementation MPPImageSegmenterTests
+
+#pragma mark General Tests
+
+- (void)setUp {
+  // When expected and actual mask sizes are not equal, iterating through mask data results in a
+  // segmentation fault. Setting this property to `NO`, prevents each test case from executing the
+  // remaining flow after a failure. Since expected and actual mask sizes are compared before
+  // iterating through them, this prevents any illegal memory access.
+  self.continueAfterFailure = NO;
+}
+
++ (NSString *)filePathWithName:(NSString *)fileName extension:(NSString *)extension {
+  NSString *filePath =
+      [[NSBundle bundleForClass:[MPPImageSegmenterTests class]] pathForResource:fileName
+                                                                         ofType:extension];
+  return filePath;
+}
+
+#pragma mark Image Mode Tests
+
+- (void)testSegmentWithCategoryMaskSucceeds {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModel];
+  options.shouldOutputConfidenceMasks = NO;
+  options.shouldOutputCategoryMask = YES;
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  [self assertResultsOfSegmentImageWithFileInfo:kSegmentationImageFileInfo
+                                           usingImageSegmenter:imageSegmenter
+      approximatelyEqualsExpectedCategoryMaskImageWithFileInfo:kSegmentationGoldenImageFileInfo
+                                     shouldHaveConfidenceMasks:NO];
+}
+
+- (void)testSegmentWithConfidenceMaskSucceeds {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModel];
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  [self assertResultsOfSegmentImageWithFileInfo:kCatImageFileInfo
+                                             usingImageSegmenter:imageSegmenter
+      approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:kCatGoldenImageFileInfo
+                                                         atIndex:8
+                                          shouldHaveCategoryMask:NO];
+}
+
+#pragma mark - Image Segmenter Initializers
+
+- (MPPImageSegmenterOptions *)imageSegmenterOptionsWithModelFileInfo:(MPPFileInfo *)fileInfo {
+  MPPImageSegmenterOptions *options = [[MPPImageSegmenterOptions alloc] init];
+  options.baseOptions.modelAssetPath = fileInfo.path;
+  return options;
+}
+
+- (MPPImageSegmenter *)createImageSegmenterWithOptionsSucceeds:(MPPImageSegmenterOptions *)options {
+  NSError *error;
+  MPPImageSegmenter *imageSegmenter = [[MPPImageSegmenter alloc] initWithOptions:options
+                                                                           error:&error];
+  XCTAssertNotNil(imageSegmenter);
+  XCTAssertNil(error);
+
+  return imageSegmenter;
+}
+
+#pragma mark Assert Segmenter Results
+- (void)assertResultsOfSegmentImageWithFileInfo:(MPPFileInfo *)imageFileInfo
+                                         usingImageSegmenter:(MPPImageSegmenter *)imageSegmenter
+    approximatelyEqualsExpectedCategoryMaskImageWithFileInfo:
+        (MPPFileInfo *)expectedCategoryMaskFileInfo
+                                   shouldHaveConfidenceMasks:(BOOL)shouldHaveConfidenceMasks {
+  MPPImageSegmenterResult *result = [self segmentImageWithFileInfo:imageFileInfo
+                                               usingImageSegmenter:imageSegmenter];
+
+  XCTAssertNotNil(result.categoryMask);
+
+  if (shouldHaveConfidenceMasks) {
+    XCTAssertNotNil(result.confidenceMasks);
+  } else {
+    XCTAssertNil(result.confidenceMasks);
+  }
+
+  [self assertCategoryMask:result.categoryMask
+      approximatelyEqualsExpectedCategoryMaskImageWithFileInfo:expectedCategoryMaskFileInfo];
+}
+
+- (void)assertResultsOfSegmentImageWithFileInfo:(MPPFileInfo *)imageFileInfo
+                                           usingImageSegmenter:(MPPImageSegmenter *)imageSegmenter
+    approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:
+        (MPPFileInfo *)expectedConfidenceMaskFileInfo
+                                                       atIndex:(NSInteger)index
+                                        shouldHaveCategoryMask:(BOOL)shouldHaveCategoryMask {
+  MPPImageSegmenterResult *result = [self segmentImageWithFileInfo:imageFileInfo
+                                               usingImageSegmenter:imageSegmenter];
+
+  XCTAssertNotNil(result.confidenceMasks);
+
+  if (shouldHaveCategoryMask) {
+    XCTAssertNotNil(result.categoryMask);
+  } else {
+    XCTAssertNil(result.categoryMask);
+  }
+
+  XCTAssertLessThan(index, result.confidenceMasks.count);
+
+  [self assertConfidenceMask:result.confidenceMasks[index]
+      approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:expectedConfidenceMaskFileInfo];
+}
+
+- (MPPImageSegmenterResult *)segmentImageWithFileInfo:(MPPFileInfo *)fileInfo
+                                  usingImageSegmenter:(MPPImageSegmenter *)imageSegmenter {
+  MPPImage *image = [MPPImage imageWithFileInfo:fileInfo];
+  XCTAssertNotNil(image);
+
+  NSError *error;
+  MPPImageSegmenterResult *result = [imageSegmenter segmentImage:image error:&error];
+  XCTAssertNil(error);
+  XCTAssertNotNil(result);
+
+  return result;
+}
+
+- (void)assertCategoryMask:(MPPMask *)categoryMask
+    approximatelyEqualsExpectedCategoryMaskImageWithFileInfo:
+        (MPPFileInfo *)expectedCategoryMaskImageFileInfo {
+  MPPMask *expectedCategoryMask =
+      [[MPPMask alloc] initWithImageFileInfo:expectedCategoryMaskImageFileInfo];
+
+  XCTAssertEqual(categoryMask.width, expectedCategoryMask.width);
+  XCTAssertEqual(categoryMask.height, expectedCategoryMask.height);
+
+  size_t maskSize = categoryMask.width * categoryMask.height;
+
+  const UInt8 *categoryMaskPixelData = categoryMask.uint8Data;
+  const UInt8 *expectedCategoryMaskPixelData = expectedCategoryMask.uint8Data;
+
+  NSInteger consistentPixels = 0;
+
+  for (int i = 0; i < maskSize; i++) {
+    consistentPixels +=
+        categoryMaskPixelData[i] * kMagnificationFactor == expectedCategoryMaskPixelData[i] ? 1 : 0;
+  }
+
+  XCTAssertGreaterThan((float)consistentPixels / (float)maskSize, kSimilarityThreshold);
+}
+
+- (void)assertConfidenceMask:(MPPMask *)confidenceMask
+    approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:
+        (MPPFileInfo *)expectedConfidenceMaskImageFileInfo {
+  MPPMask *expectedConfidenceMask =
+      [[MPPMask alloc] initWithImageFileInfo:expectedConfidenceMaskImageFileInfo];
+
+  XCTAssertEqual(confidenceMask.width, expectedConfidenceMask.width);
+  XCTAssertEqual(confidenceMask.height, expectedConfidenceMask.height);
+
+  size_t maskSize = confidenceMask.width * confidenceMask.height;
+
+  XCTAssertGreaterThan(
+      softIOU(confidenceMask.float32Data, expectedConfidenceMask.float32Data, maskSize),
+      kSimilarityThreshold);
+}
+
+@end

mediapipe/tasks/ios/test/vision/utils/BUILD

@@ -1,3 +1,17 @@
+# Copyright 2023 The MediaPipe Authors. All Rights Reserved.
+#
+# 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.
+
 package(default_visibility = ["//mediapipe/tasks:internal"])
 
 licenses(["notice"])
@@ -7,7 +21,22 @@ 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/test/utils:MPPFileInfo",
+        "//mediapipe/tasks/ios/vision/core:MPPImage",
+    ],
+)
+
+objc_library(
+    name = "MPPMaskTestUtils",
+    srcs = ["sources/MPPMask+TestUtils.m"],
+    hdrs = ["sources/MPPMask+TestUtils.h"],
+    module_name = "MPPMaskTestUtils",
+    deps = [
+        "//mediapipe/tasks/ios/test/utils:MPPFileInfo",
+        "//mediapipe/tasks/ios/vision/core:MPPMask",
+        "//third_party/apple_frameworks:UIKit",
+    ],
 )
 
 cc_library(

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,31 @@ 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.
  *
@@ -31,14 +57,15 @@ NS_ASSUME_NONNULL_BEGIN
  * @param name Name of the image file.
  * @param type Extension of the image file.
  *
- * @return The `MPPImage` object contains the loaded image. This method returns
- * nil if it cannot load the image.
+ * @return The `MPPImage` object contains the loaded image. This method returns nil if it cannot
+ * load the image.
  */
 + (nullable MPPImage *)imageFromBundleWithClass:(Class)classObject
                                        fileName:(NSString *)name
                                          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.
@@ -49,8 +76,8 @@ NS_ASSUME_NONNULL_BEGIN
  * @param type Extension of the image file.
  * @param orientation Orientation of the image.
  *
- * @return The `MPPImage` object contains the loaded image. This method returns
- * nil if it cannot load the image.
+ * @return The `MPPImage` object contains the loaded image. This method returns nil if it cannot
+ * load the image.
  */
 + (nullable MPPImage *)imageFromBundleWithClass:(Class)classObject
                                        fileName:(NSString *)name

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

mediapipe/tasks/ios/test/vision/utils/sources/MPPMask+TestUtils.h

@@ -0,0 +1,42 @@
+// 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 <Foundation/Foundation.h>
+#import <UIKit/UIKit.h>
+
+#import "mediapipe/tasks/ios/test/utils/sources/MPPFileInfo.h"
+#import "mediapipe/tasks/ios/vision/core/sources/MPPMask.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+/**
+ * Helper utility for initializing `MPPMask` for MediaPipe iOS vision library tests.
+ */
+@interface MPPMask (TestUtils)
+
+/**
+ * Loads an image from a file in an app bundle and Creates an `MPPMask` of type
+ * `MPPMaskDataTypeUInt8` using the gray scale pixel data of a `UIImage` loaded from a file with the
+ * given `MPPFileInfo`.
+ *
+ * @param fileInfo The file info specifying the name and type of the image file in the app bundle.
+ *
+ * @return The `MPPMask` with the pixel data of the loaded image. This method returns `nil` if there
+ * is an error in loading the image correctly.
+ */
+- (nullable instancetype)initWithImageFileInfo:(MPPFileInfo *)fileInfo;
+
+@end
+
+NS_ASSUME_NONNULL_END

mediapipe/tasks/ios/test/vision/utils/sources/MPPMask+TestUtils.m

@@ -0,0 +1,60 @@
+// 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/test/vision/utils/sources/MPPMask+TestUtils.h"
+
+@implementation MPPMask (TestUtils)
+
+- (instancetype)initWithImageFileInfo:(MPPFileInfo *)fileInfo {
+  UIImage *image = [[UIImage alloc] initWithContentsOfFile:fileInfo.path];
+
+  if (!image.CGImage) {
+    return nil;
+  }
+
+  size_t width = CGImageGetWidth(image.CGImage);
+  size_t height = CGImageGetHeight(image.CGImage);
+
+  NSInteger bitsPerComponent = 8;
+
+  UInt8 *pixelData = NULL;
+
+  CGColorSpaceRef colorSpace = CGColorSpaceCreateDeviceGray();
+
+  // For a gray scale image (single component) with no alpha, the bitmap info is
+  // `kCGImageAlphaNone` in combination with bytesPerRow being equal to width.
+  CGContextRef context = CGBitmapContextCreate(nil, width, height, bitsPerComponent, width,
+                                               colorSpace, kCGImageAlphaNone);
+
+  if (!context) {
+    CGColorSpaceRelease(colorSpace);
+    return nil;
+  }
+
+  CGContextDrawImage(context, CGRectMake(0, 0, width, height), image.CGImage);
+  pixelData = (UInt8 *)CGBitmapContextGetData(context);
+
+  // A copy is needed to ensure that the pixel data outlives the `CGContextRelease` call.
+  // Alternative is to make the context, color space instance variables and release them in
+  // `dealloc()`. Since Categories don't allow adding instance variables, choosing to copy rather
+  // than creating a new custom class similar to `MPPMask` only for the tests.
+  MPPMask *mask = [self initWithUInt8Data:pixelData width:width height:height shouldCopy:YES];
+
+  CGColorSpaceRelease(colorSpace);
+  CGContextRelease(context);
+
+  return mask;
+}
+
+@end

mediapipe/tasks/ios/vision/image_segmenter/utils/sources/MPPImageSegmenterResult+Helpers.mm

@@ -52,7 +52,7 @@ using ::mediapipe::Packet;
   }
 
   if (categoryMaskPacket.ValidateAsType<Image>().ok()) {
-    const Image &cppCategoryMask = confidenceMasksPacket.Get<Image>();
+    const Image &cppCategoryMask = categoryMaskPacket.Get<Image>();
     categoryMask = [[MPPMask alloc]
         initWithUInt8Data:(UInt8 *)cppCategoryMask.GetImageFrameSharedPtr().get()->PixelData()
                     width:cppCategoryMask.width()