Merge pull request #4830 from priankakariatyml:ios-image-segmenter-ba…

…sic-tests PiperOrigin-RevId: 571386542
google-ai-edge · Oct 6, 2023 · 830ee09 · 830ee09
2 parents 90e6a97 + b1f717e
commit 830ee09
Showing 1 changed file with 39 additions and 0 deletions.
diff --git a/mediapipe/tasks/ios/test/vision/image_segmenter/MPPImageSegmenterTests.mm b/mediapipe/tasks/ios/test/vision/image_segmenter/MPPImageSegmenterTests.mm
@@ -468,6 +468,45 @@ - (void)imageSegmenter:(MPPImageSegmenter *)imageSegmenter
   }
 }
 
+#pragma mark Mask No Copy Tests
+
+- (void)testSegmentWithNoCopyConfidenceMasksAndImageModeSucceeds {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  MPPImage *image = [MPPImage imageWithFileInfo:kCatImageFileInfo];
+  [imageSegmenter segmentImage:image withCompletionHandler:^(MPPImageSegmenterResult *result, NSError *error) {
+      [self assertImageSegmenterResult:result
+                                         hasConfidenceMasksCount:
+                                             kExpectedDeeplabV3ConfidenceMaskCount
+      approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:kCatGoldenImageFileInfo
+                                                         atIndex:8
+                                          shouldHaveCategoryMask:NO];
+  }];
+}
+
+- (void)testSegmentWithNoCopyConfidenceMasksAndVideoModeSucceeds {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
+  options.runningMode = MPPRunningModeVideo;
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  const NSInteger timestampInMilliseconds = 0;
+
+  MPPImage *image = [MPPImage imageWithFileInfo:kCatImageFileInfo];
+  [imageSegmenter segmentVideoFrame:image timestampInMilliseconds:timestampInMilliseconds withCompletionHandler:^(MPPImageSegmenterResult *result, NSError *error) {
+      [self assertImageSegmenterResult:result
+                                         hasConfidenceMasksCount:
+                                             kExpectedDeeplabV3ConfidenceMaskCount
+      approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:kCatGoldenImageFileInfo
+                                                         atIndex:8
+                                          shouldHaveCategoryMask:NO];
+  }];
+}
+
 #pragma mark - Image Segmenter Initializers
 
 - (MPPImageSegmenterOptions *)imageSegmenterOptionsWithModelFileInfo:(MPPFileInfo *)fileInfo {