Added selfie segmentation and running mode tests to image segmenter

This commit is contained in:
Prianka Liz Kariat 2023-09-19 20:00:41 +05:30
parent 08a5d55ac1
commit bac60548dc
2 changed files with 382 additions and 6 deletions

View File

@ -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",

View File

@ -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) {
@ -72,7 +98,10 @@ double softIOU(const float *mask1, const float *mask2, size_t size) {
}
} // 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 {