Added selfie segmentation and running mode tests to image segmenter

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

@@ -53,6 +53,7 @@ objc_library(
         "//mediapipe/tasks/testdata/vision:test_protos",
     ],
     deps = [
+        "//mediapipe/tasks/ios/common:MPPCommon",
         "//mediapipe/tasks/ios/test/vision/utils:MPPImageTestUtils",
         "//mediapipe/tasks/ios/test/vision/utils:MPPMaskTestUtils",
         "//mediapipe/tasks/ios/vision/image_segmenter:MPPImageSegmenter",

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

@@ -15,6 +15,7 @@
 #import <Foundation/Foundation.h>
 #import <XCTest/XCTest.h>
 
+#import "mediapipe/tasks/ios/common/sources/MPPCommon.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"
@@ -30,11 +31,36 @@ 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 MPPFileInfo *const kMozartImageFileInfo = [[MPPFileInfo alloc] initWithName:@"mozart_square"
+                                                                              type:@"jpg"];
+static MPPFileInfo *const kMozart128x128SegmentationGoldenImageFileInfo =
+    [[MPPFileInfo alloc] initWithName:@"selfie_segm_128_128_3_expected_mask" type:@"jpg"];
+static MPPFileInfo *const kMozart144x256SegmentationGoldenImageFileInfo =
+    [[MPPFileInfo alloc] initWithName:@"selfie_segm_144_256_3_expected_mask" type:@"jpg"];
+
+static MPPFileInfo *const kImageSegmenterModelFileInfo =
+    [[MPPFileInfo alloc] initWithName:@"deeplabv3" type:@"tflite"];
+static MPPFileInfo *const kSelfie128x128ModelFileInfo =
+    [[MPPFileInfo alloc] initWithName:@"selfie_segm_128_128_3" type:@"tflite"];
+static MPPFileInfo *const kSelfie144x256ModelFileInfo =
+    [[MPPFileInfo alloc] initWithName:@"selfie_segm_144_256_3" type:@"tflite"];
+
 static NSString *const kExpectedErrorDomain = @"com.google.mediapipe.tasks";
+static NSString *const kLiveStreamTestsDictImageSegmenterKey = @"image_segmenter";
+static NSString *const kLiveStreamTestsDictExpectationKey = @"expectation";
+
 constexpr float kSimilarityThreshold = 0.96f;
 constexpr NSInteger kMagnificationFactor = 10;
+constexpr NSInteger kExpectedDeeplabV3ConfidenceMaskCount = 21;
+constexpr NSInteger kExpected128x128SelfieSegmentationConfidenceMaskCount = 2;
+constexpr NSInteger kExpected144x256SelfieSegmentationConfidenceMaskCount = 1;
+
+#define AssertEqualErrors(error, expectedError)              \
+  XCTAssertNotNil(error);                                    \
+  XCTAssertEqualObjects(error.domain, expectedError.domain); \
+  XCTAssertEqual(error.code, expectedError.code);            \
+  XCTAssertEqualObjects(error.localizedDescription, expectedError.localizedDescription)
 
 namespace {
 double sum(const std::vector<float> &mask) {
@@ -70,9 +96,12 @@ double softIOU(const float *mask1, const float *mask2, size_t size) {
 
   return unionSum > 0.0 ? interSectionSum / unionSum : 0.0;
 }
-} // namespace
+}  // namespace
 
-@interface MPPImageSegmenterTests : XCTestCase <MPPImageSegmenterLiveStreamDelegate>
+@interface MPPImageSegmenterTests : XCTestCase <MPPImageSegmenterLiveStreamDelegate> {
+  NSDictionary<NSString *, id> *_liveStreamSucceedsTestDict;
+  NSDictionary<NSString *, id> *_outOfOrderTimestampTestDict;
+}
 
 @end
 
@@ -99,7 +128,7 @@ double softIOU(const float *mask1, const float *mask2, size_t size) {
 
 - (void)testSegmentWithCategoryMaskSucceeds {
   MPPImageSegmenterOptions *options =
-      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModel];
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
   options.shouldOutputConfidenceMasks = NO;
   options.shouldOutputCategoryMask = YES;
 
@@ -113,17 +142,335 @@ double softIOU(const float *mask1, const float *mask2, size_t size) {
 
 - (void)testSegmentWithConfidenceMaskSucceeds {
   MPPImageSegmenterOptions *options =
-      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModel];
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
 
   MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
 
   [self assertResultsOfSegmentImageWithFileInfo:kCatImageFileInfo
                                              usingImageSegmenter:imageSegmenter
+                                         hasConfidenceMasksCount:
+                                             kExpectedDeeplabV3ConfidenceMaskCount
       approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:kCatGoldenImageFileInfo
                                                          atIndex:8
                                           shouldHaveCategoryMask:NO];
 }
 
+- (void)testSegmentWith128x128SegmentationSucceeds {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kSelfie128x128ModelFileInfo];
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  [self assertResultsOfSegmentImageWithFileInfo:kMozartImageFileInfo
+                                             usingImageSegmenter:imageSegmenter
+                                         hasConfidenceMasksCount:
+                                             kExpected128x128SelfieSegmentationConfidenceMaskCount
+      approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:
+          kMozart128x128SegmentationGoldenImageFileInfo
+                                                         atIndex:1
+                                          shouldHaveCategoryMask:NO];
+}
+
+- (void)testSegmentWith144x256SegmentationSucceeds {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kSelfie144x256ModelFileInfo];
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  [self assertResultsOfSegmentImageWithFileInfo:kMozartImageFileInfo
+                                             usingImageSegmenter:imageSegmenter
+                                         hasConfidenceMasksCount:
+                                             kExpected144x256SelfieSegmentationConfidenceMaskCount
+      approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:
+          kMozart144x256SegmentationGoldenImageFileInfo
+                                                         atIndex:0
+                                          shouldHaveCategoryMask:NO];
+}
+
+#pragma mark Running Mode Tests
+
+- (void)testCreateImageSegmenterFailsWithDelegateInNonLiveStreamMode {
+  MPPRunningMode runningModesToTest[] = {MPPRunningModeImage, MPPRunningModeVideo};
+  for (int i = 0; i < sizeof(runningModesToTest) / sizeof(runningModesToTest[0]); i++) {
+    MPPImageSegmenterOptions *options =
+        [self imageSegmenterOptionsWithModelFileInfo:kSelfie128x128ModelFileInfo];
+
+    options.runningMode = runningModesToTest[i];
+    options.imageSegmenterLiveStreamDelegate = self;
+
+    [self
+        assertCreateImageSegmenterWithOptions:options
+                       failsWithExpectedError:
+                           [NSError errorWithDomain:kExpectedErrorDomain
+                                               code:MPPTasksErrorCodeInvalidArgumentError
+                                           userInfo:@{
+                                             NSLocalizedDescriptionKey :
+                                                 @"The vision task is in image or video mode. The "
+                                                 @"delegate must not be set in the task's options."
+                                           }]];
+  }
+}
+
+- (void)testCreateImageSegmenterFailsWithMissingDelegateInLiveStreamMode {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kSelfie128x128ModelFileInfo];
+
+  options.runningMode = MPPRunningModeLiveStream;
+
+  [self assertCreateImageSegmenterWithOptions:options
+                       failsWithExpectedError:
+                           [NSError
+                               errorWithDomain:kExpectedErrorDomain
+                                          code:MPPTasksErrorCodeInvalidArgumentError
+                                      userInfo:@{
+                                        NSLocalizedDescriptionKey :
+                                            @"The vision task is in live stream mode. An object "
+                                            @"must be set as the delegate of the task in its "
+                                            @"options to ensure asynchronous delivery of results."
+                                      }]];
+}
+
+- (void)testSegmentFailsWithCallingWrongApiInImageMode {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  MPPImage *image = [MPPImage imageWithFileInfo:kCatImageFileInfo];
+  XCTAssertNotNil(image);
+
+  NSError *liveStreamApiCallError;
+  XCTAssertFalse([imageSegmenter segmentAsyncImage:image
+                           timestampInMilliseconds:0
+                                             error:&liveStreamApiCallError]);
+
+  NSError *expectedLiveStreamApiCallError =
+      [NSError errorWithDomain:kExpectedErrorDomain
+                          code:MPPTasksErrorCodeInvalidArgumentError
+                      userInfo:@{
+                        NSLocalizedDescriptionKey : @"The vision task is not initialized with live "
+                                                    @"stream mode. Current Running Mode: Image"
+                      }];
+
+  AssertEqualErrors(liveStreamApiCallError, expectedLiveStreamApiCallError);
+
+  NSError *videoApiCallError;
+  XCTAssertFalse([imageSegmenter segmentVideoFrame:image
+                           timestampInMilliseconds:0
+                                             error:&videoApiCallError]);
+
+  NSError *expectedVideoApiCallError =
+      [NSError errorWithDomain:kExpectedErrorDomain
+                          code:MPPTasksErrorCodeInvalidArgumentError
+                      userInfo:@{
+                        NSLocalizedDescriptionKey : @"The vision task is not initialized with "
+                                                    @"video mode. Current Running Mode: Image"
+                      }];
+  AssertEqualErrors(videoApiCallError, expectedVideoApiCallError);
+}
+
+- (void)testSegmentFailsWithCallingWrongApiInVideoMode {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
+  options.runningMode = MPPRunningModeVideo;
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  MPPImage *image = [MPPImage imageWithFileInfo:kCatImageFileInfo];
+  XCTAssertNotNil(image);
+
+  NSError *liveStreamApiCallError;
+  XCTAssertFalse([imageSegmenter segmentAsyncImage:image
+                           timestampInMilliseconds:0
+                                             error:&liveStreamApiCallError]);
+
+  NSError *expectedLiveStreamApiCallError =
+      [NSError errorWithDomain:kExpectedErrorDomain
+                          code:MPPTasksErrorCodeInvalidArgumentError
+                      userInfo:@{
+                        NSLocalizedDescriptionKey : @"The vision task is not initialized with live "
+                                                    @"stream mode. Current Running Mode: Video"
+                      }];
+
+  AssertEqualErrors(liveStreamApiCallError, expectedLiveStreamApiCallError);
+
+  NSError *imageApiCallError;
+  XCTAssertFalse([imageSegmenter segmentImage:image error:&imageApiCallError]);
+
+  NSError *expectedImageApiCallError =
+      [NSError errorWithDomain:kExpectedErrorDomain
+                          code:MPPTasksErrorCodeInvalidArgumentError
+                      userInfo:@{
+                        NSLocalizedDescriptionKey : @"The vision task is not initialized with "
+                                                    @"image mode. Current Running Mode: Video"
+                      }];
+  AssertEqualErrors(imageApiCallError, expectedImageApiCallError);
+}
+
+- (void)testSegmentFailsWithCallingWrongApiInLiveStreamMode {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
+  options.runningMode = MPPRunningModeLiveStream;
+  options.imageSegmenterLiveStreamDelegate = self;
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  MPPImage *image = [MPPImage imageWithFileInfo:kCatImageFileInfo];
+  XCTAssertNotNil(image);
+
+  NSError *imageApiCallError;
+  XCTAssertFalse([imageSegmenter segmentImage:image error:&imageApiCallError]);
+
+  NSError *expectedImageApiCallError =
+      [NSError errorWithDomain:kExpectedErrorDomain
+                          code:MPPTasksErrorCodeInvalidArgumentError
+                      userInfo:@{
+                        NSLocalizedDescriptionKey : @"The vision task is not initialized with "
+                                                    @"image mode. Current Running Mode: Live Stream"
+                      }];
+  AssertEqualErrors(imageApiCallError, expectedImageApiCallError);
+
+  NSError *videoApiCallError;
+  XCTAssertFalse([imageSegmenter segmentVideoFrame:image
+                           timestampInMilliseconds:0
+                                             error:&videoApiCallError]);
+
+  NSError *expectedVideoApiCallError =
+      [NSError errorWithDomain:kExpectedErrorDomain
+                          code:MPPTasksErrorCodeInvalidArgumentError
+                      userInfo:@{
+                        NSLocalizedDescriptionKey : @"The vision task is not initialized with "
+                                                    @"video mode. Current Running Mode: Live Stream"
+                      }];
+  AssertEqualErrors(videoApiCallError, expectedVideoApiCallError);
+}
+
+- (void)testSegmentWithVideoModeSucceeds {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
+  options.runningMode = MPPRunningModeVideo;
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  MPPImage *image = [MPPImage imageWithFileInfo:kCatImageFileInfo];
+  XCTAssertNotNil(image);
+
+  for (int i = 0; i < 3; i++) {
+    MPPImageSegmenterResult *result = [imageSegmenter segmentVideoFrame:image
+                                                timestampInMilliseconds:i
+                                                                  error:nil];
+    [self assertImageSegmenterResult:result
+                                           hasConfidenceMasksCount:
+                                               kExpectedDeeplabV3ConfidenceMaskCount
+        approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:kCatGoldenImageFileInfo
+                                                           atIndex:8
+                                            shouldHaveCategoryMask:NO];
+  }
+}
+
+// - (void)testSegmentWithOutOfOrderTimestampsAndLiveStreamModeFails {
+//  MPPImageSegmenterOptions *options =
+//       [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
+//   options.runningMode = MPPRunningModeLiveStream;
+//   options.imageSegmenterLiveStreamDelegate = self;
+
+//   XCTestExpectation *expectation = [[XCTestExpectation alloc]
+//       initWithDescription:@"segmentWithOutOfOrderTimestampsAndLiveStream"];
+
+//   expectation.expectedFulfillmentCount = 1;
+
+//   MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+//   _outOfOrderTimestampTestDict = @{
+//     kLiveStreamTestsDictImageSegmenterKey : imageSegmenter,
+//     kLiveStreamTestsDictExpectationKey : expectation
+//   };
+
+//   MPPImage *image = [MPPImage imageWithFileInfo:kCatImageFileInfo];
+//   XCTAssertNotNil(image);
+
+//   XCTAssertTrue([imageSegmenter segmentAsyncImage:image timestampInMilliseconds:1 error:nil]);
+
+//   NSError *error;
+//   XCTAssertFalse([imageSegmenter segmentAsyncImage:image
+//                                 timestampInMilliseconds:0
+//                                                   error:&error]);
+
+//   NSError *expectedError =
+//       [NSError errorWithDomain:kExpectedErrorDomain
+//                           code:MPPTasksErrorCodeInvalidArgumentError
+//                       userInfo:@{
+//                         NSLocalizedDescriptionKey :
+//                             @"INVALID_ARGUMENT: Input timestamp must be monotonically
+//                             increasing."
+//                       }];
+//   AssertEqualErrors(error, expectedError);
+
+//   NSTimeInterval timeout = 0.5f;
+//   [self waitForExpectations:@[ expectation ] timeout:timeout];
+// }
+
+- (void)testSegmentWithLiveStreamModeSucceeds {
+  MPPImageSegmenterOptions *options =
+      [self imageSegmenterOptionsWithModelFileInfo:kImageSegmenterModelFileInfo];
+  options.runningMode = MPPRunningModeLiveStream;
+  options.imageSegmenterLiveStreamDelegate = self;
+
+  NSInteger iterationCount = 100;
+
+  // Because of flow limiting, we cannot ensure that the callback will be invoked `iterationCount`
+  // times. An normal expectation will fail if expectation.fulfill() is not called
+  // `expectation.expectedFulfillmentCount` times. If `expectation.isInverted = true`, the test will
+  // only succeed if expectation is not fulfilled for the specified `expectedFulfillmentCount`.
+  // Since in our case we cannot predict how many times the expectation is supposed to be fulfilled
+  // setting, `expectation.expectedFulfillmentCount` = `iterationCount` + 1 and
+  // `expectation.isInverted = true` ensures that test succeeds ifexpectation is fulfilled <=
+  // `iterationCount` times.
+  XCTestExpectation *expectation =
+      [[XCTestExpectation alloc] initWithDescription:@"segmentWithLiveStream"];
+
+  expectation.expectedFulfillmentCount = iterationCount + 1;
+  expectation.inverted = YES;
+
+  MPPImageSegmenter *imageSegmenter = [self createImageSegmenterWithOptionsSucceeds:options];
+
+  _outOfOrderTimestampTestDict = @{
+    kLiveStreamTestsDictImageSegmenterKey : imageSegmenter,
+    kLiveStreamTestsDictExpectationKey : expectation
+  };
+
+  // TODO: Mimic initialization from CMSampleBuffer as live stream mode is most likely to be used
+  // with the iOS camera. AVCaptureVideoDataOutput sample buffer delegates provide frames of type
+  // `CMSampleBuffer`.
+  MPPImage *image = [MPPImage imageWithFileInfo:kCatImageFileInfo];
+  XCTAssertNotNil(image);
+
+  for (int i = 0; i < iterationCount; i++) {
+    XCTAssertTrue([imageSegmenter segmentAsyncImage:image timestampInMilliseconds:i error:nil]);
+  }
+
+  NSTimeInterval timeout = 0.5f;
+  [self waitForExpectations:@[ expectation ] timeout:timeout];
+}
+
+- (void)imageSegmenter:(MPPImageSegmenter *)imageSegmenter
+    didFinishSegmentationWithResult:(MPPImageSegmenterResult *)imageSegmenterResult
+            timestampInMilliseconds:(NSInteger)timestampInMilliseconds
+                              error:(NSError *)error {
+  [self assertImageSegmenterResult:imageSegmenterResult
+                                         hasConfidenceMasksCount:
+                                             kExpectedDeeplabV3ConfidenceMaskCount
+      approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:kCatGoldenImageFileInfo
+                                                         atIndex:8
+                                          shouldHaveCategoryMask:NO];
+
+  if (imageSegmenter == _outOfOrderTimestampTestDict[kLiveStreamTestsDictImageSegmenterKey]) {
+    [_outOfOrderTimestampTestDict[kLiveStreamTestsDictExpectationKey] fulfill];
+  } else if (imageSegmenter == _liveStreamSucceedsTestDict[kLiveStreamTestsDictImageSegmenterKey]) {
+    [_liveStreamSucceedsTestDict[kLiveStreamTestsDictExpectationKey] fulfill];
+  }
+}
+
 #pragma mark - Image Segmenter Initializers
 
 - (MPPImageSegmenterOptions *)imageSegmenterOptionsWithModelFileInfo:(MPPFileInfo *)fileInfo {
@@ -142,6 +489,16 @@ double softIOU(const float *mask1, const float *mask2, size_t size) {
   return imageSegmenter;
 }
 
+- (void)assertCreateImageSegmenterWithOptions:(MPPImageSegmenterOptions *)options
+                       failsWithExpectedError:(NSError *)expectedError {
+  NSError *error = nil;
+  MPPImageSegmenter *imageSegmenter = [[MPPImageSegmenter alloc] initWithOptions:options
+                                                                           error:&error];
+
+  XCTAssertNil(imageSegmenter);
+  AssertEqualErrors(error, expectedError);
+}
+
 #pragma mark Assert Segmenter Results
 - (void)assertResultsOfSegmentImageWithFileInfo:(MPPFileInfo *)imageFileInfo
                                          usingImageSegmenter:(MPPImageSegmenter *)imageSegmenter
@@ -165,6 +522,8 @@ double softIOU(const float *mask1, const float *mask2, size_t size) {
 
 - (void)assertResultsOfSegmentImageWithFileInfo:(MPPFileInfo *)imageFileInfo
                                            usingImageSegmenter:(MPPImageSegmenter *)imageSegmenter
+                                       hasConfidenceMasksCount:
+                                           (NSUInteger)expectedConfidenceMasksCount
     approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:
         (MPPFileInfo *)expectedConfidenceMaskFileInfo
                                                        atIndex:(NSInteger)index
@@ -172,8 +531,24 @@ double softIOU(const float *mask1, const float *mask2, size_t size) {
   MPPImageSegmenterResult *result = [self segmentImageWithFileInfo:imageFileInfo
                                                usingImageSegmenter:imageSegmenter];
 
+  [self assertImageSegmenterResult:result
+                                         hasConfidenceMasksCount:expectedConfidenceMasksCount
+      approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:expectedConfidenceMaskFileInfo
+                                                         atIndex:index
+                                          shouldHaveCategoryMask:shouldHaveCategoryMask];
+}
+
+- (void)assertImageSegmenterResult:(MPPImageSegmenterResult *)result
+                                       hasConfidenceMasksCount:
+                                           (NSUInteger)expectedConfidenceMasksCount
+    approximatelyEqualsExpectedConfidenceMaskImageWithFileInfo:
+        (MPPFileInfo *)expectedConfidenceMaskFileInfo
+                                                       atIndex:(NSInteger)index
+                                        shouldHaveCategoryMask:(BOOL)shouldHaveCategoryMask {
   XCTAssertNotNil(result.confidenceMasks);
 
+  XCTAssertEqual(result.confidenceMasks.count, expectedConfidenceMasksCount);
+
   if (shouldHaveCategoryMask) {
     XCTAssertNotNil(result.categoryMask);
   } else {