Added more utility methods to MPPAudioEmbedderTests

kp-forks · Oct 11, 2024 · fcb77b5 · fcb77b5
1 parent c9851d9
commit fcb77b5
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 60 deletions.
diff --git a/mediapipe/tasks/ios/test/audio/audio_classifier/MPPAudioClassifierTests.mm b/mediapipe/tasks/ios/test/audio/audio_classifier/MPPAudioClassifierTests.mm
@@ -352,8 +352,7 @@ - (void)testClassifyFailsWithCallingWrongApiInAudioClipsMode {
   MPPAudioClassifier *audioClassifier =
       [MPPAudioClassifierTests audioClassifierWithOptions:options];
 
-  MPPAudioData *audioClip =
-     [[MPPAudioData alloc] initWithFileInfo:kSpeech16KHzMonoFileInfo];
+  MPPAudioData *audioClip = [[MPPAudioData alloc] initWithFileInfo:kSpeech16KHzMonoFileInfo];
   NSError *error;
   XCTAssertFalse([audioClassifier classifyAsyncAudioBlock:audioClip
                                   timestampInMilliseconds:0
@@ -380,8 +379,7 @@ - (void)testClassifyFailsWithCallingWrongApiInAudioStreamMode {
   MPPAudioClassifier *audioClassifier =
       [MPPAudioClassifierTests audioClassifierWithOptions:options];
 
-  MPPAudioData *audioClip =
-      [[MPPAudioData alloc] initWithFileInfo:kSpeech16KHzMonoFileInfo];
+  MPPAudioData *audioClip = [[MPPAudioData alloc] initWithFileInfo:kSpeech16KHzMonoFileInfo];
 
   NSError *error;
   XCTAssertFalse([audioClassifier classifyAudioClip:audioClip error:&error]);
@@ -442,9 +440,9 @@ - (void)testClassifyWithAudioStreamModeAndOutOfOrderTimestampsFails {
 
 - (void)testClassifyWithAudioStreamModeSucceeds {
   [self classifyUsingYamnetAsyncAudioFileWithInfo:kSpeech16KHzMonoFileInfo
-                                                 info:&_16kHZAudioStreamSucceedsTestDict];
+                                             info:&_16kHZAudioStreamSucceedsTestDict];
   [self classifyUsingYamnetAsyncAudioFileWithInfo:kSpeech48KHzMonoFileInfo
-                                                 info:&_48kHZAudioStreamSucceedsTestDict];
+                                             info:&_48kHZAudioStreamSucceedsTestDict];
 }
 
 - (void)audioClassifier:(MPPAudioClassifier *)audioClassifier
@@ -484,8 +482,7 @@ - (void)audioClassifier:(MPPAudioClassifier *)audioClassifier
 // info is strong here since address of global variables will be passed to this function. By default
 // `NSDictionary **` will be `NSDictionary * __autoreleasing *.
 - (void)classifyUsingYamnetAsyncAudioFileWithInfo:(MPPFileInfo *)audioFileInfo
-                                                 info:(NSDictionary<NSString *, id> *__strong *)
-                                                          info {
+                                             info:(NSDictionary<NSString *, id> *__strong *)info {
   MPPAudioClassifier *audioClassifier =
       [self audioClassifierInStreamModeWithModelFileInfo:kYamnetModelFileInfo];
 
@@ -532,7 +529,7 @@ - (void)classifyUsingYamnetAsyncAudioFileWithInfo:(MPPFileInfo *)audioFileInfo
 
 - (MPPAudioClassifier *)audioClassifierInStreamModeWithModelFileInfo:(MPPFileInfo *)fileInfo {
   MPPAudioClassifierOptions *options =
-      [MPPAudioClassifierTests audioClassifierOptionsWithModelFileInfo:kYamnetModelFileInfo];
+      [MPPAudioClassifierTests audioClassifierOptionsWithModelFileInfo:fileInfo];
   options.runningMode = MPPAudioRunningModeAudioStream;
   options.audioClassifierStreamDelegate = self;
 

diff --git a/mediapipe/tasks/ios/test/audio/audio_embedder/MPPAudioEmbedderTests.mm b/mediapipe/tasks/ios/test/audio/audio_embedder/MPPAudioEmbedderTests.mm
@@ -48,7 +48,7 @@
 
 #define AssertEmbeddingResultHasOneEmbedding(embeddingResult) \
   XCTAssertNotNil(embeddingResult);                           \
-  \                                                          
+  \                                                         
   XCTAssertEqual(embeddingResult.embeddings.count, 1);
 
 #define AssertEmbeddingHasCorrectTypeAndDimension(embedding, quantize, expectedLength) \
@@ -94,21 +94,21 @@ - (void)testEmbedWithModelPathAndDifferentAudioFilesSucceeds {
   XCTAssertNotNil(audioEmbedder);
 
   [MPPAudioEmbedderTests
-      assertResultsOfEmbedAudioClipWithFileInfo:kSpeech16KHzMonoFileInfo
-                             usingAudioEmbedder:audioEmbedder
-                                    isQuantized:NO
-                  expectedEmbeddingResultsCount:kExpectedEmbeddingResultsCountForSpeechFiles];
+      assertResultsOfNonQuantizedEmbedAudioClipWithFileInfo:kSpeech16KHzMonoFileInfo
+                                         usingAudioEmbedder:audioEmbedder
+                              expectedEmbeddingResultsCount:
+                                  kExpectedEmbeddingResultsCountForSpeechFiles];
   [MPPAudioEmbedderTests
-      assertResultsOfEmbedAudioClipWithFileInfo:kSpeech48KHzMonoFileInfo
-                             usingAudioEmbedder:audioEmbedder
-                                    isQuantized:NO
-                  expectedEmbeddingResultsCount:kExpectedEmbeddingResultsCountForSpeechFiles];
+      assertResultsOfNonQuantizedEmbedAudioClipWithFileInfo:kSpeech48KHzMonoFileInfo
+                                         usingAudioEmbedder:audioEmbedder
+                              expectedEmbeddingResultsCount:
+                                  kExpectedEmbeddingResultsCountForSpeechFiles];
 
   const NSInteger expectedEmbeddingResultCount = 1;
-  [MPPAudioEmbedderTests assertResultsOfEmbedAudioClipWithFileInfo:kTwoHeads16KHzMonoFileInfo
-                                                usingAudioEmbedder:audioEmbedder
-                                                       isQuantized:NO
-                                     expectedEmbeddingResultsCount:expectedEmbeddingResultCount];
+  [MPPAudioEmbedderTests
+      assertResultsOfNonQuantizedEmbedAudioClipWithFileInfo:kTwoHeads16KHzMonoFileInfo
+                                         usingAudioEmbedder:audioEmbedder
+                              expectedEmbeddingResultsCount:expectedEmbeddingResultCount];
 }
 
 - (void)testEmbedWithOptionsSucceeds {
@@ -118,15 +118,15 @@ - (void)testEmbedWithOptionsSucceeds {
   MPPAudioEmbedder *audioEmbedder = [MPPAudioEmbedderTests audioEmbedderWithOptions:options];
 
   [MPPAudioEmbedderTests
-      assertResultsOfEmbedAudioClipWithFileInfo:kSpeech16KHzMonoFileInfo
-                             usingAudioEmbedder:audioEmbedder
-                                    isQuantized:options.quantize
-                  expectedEmbeddingResultsCount:kExpectedEmbeddingResultsCountForSpeechFiles];
+      assertResultsOfNonQuantizedEmbedAudioClipWithFileInfo:kSpeech16KHzMonoFileInfo
+                                         usingAudioEmbedder:audioEmbedder
+                              expectedEmbeddingResultsCount:
+                                  kExpectedEmbeddingResultsCountForSpeechFiles];
   [MPPAudioEmbedderTests
-      assertResultsOfEmbedAudioClipWithFileInfo:kSpeech48KHzMonoFileInfo
-                             usingAudioEmbedder:audioEmbedder
-                                    isQuantized:options.quantize
-                  expectedEmbeddingResultsCount:kExpectedEmbeddingResultsCountForSpeechFiles];
+      assertResultsOfNonQuantizedEmbedAudioClipWithFileInfo:kSpeech48KHzMonoFileInfo
+                                         usingAudioEmbedder:audioEmbedder
+                              expectedEmbeddingResultsCount:
+                                  kExpectedEmbeddingResultsCountForSpeechFiles];
 }
 
 - (void)testEmbedWithQuantizationSucceeds {
@@ -155,15 +155,15 @@ - (void)testEmbedWithL2NormalizationSucceeds {
   MPPAudioEmbedder *audioEmbedder = [MPPAudioEmbedderTests audioEmbedderWithOptions:options];
 
   [MPPAudioEmbedderTests
-      assertResultsOfEmbedAudioClipWithFileInfo:kSpeech16KHzMonoFileInfo
-                             usingAudioEmbedder:audioEmbedder
-                                    isQuantized:options.quantize
-                  expectedEmbeddingResultsCount:kExpectedEmbeddingResultsCountForSpeechFiles];
+      assertResultsOfNonQuantizedEmbedAudioClipWithFileInfo:kSpeech16KHzMonoFileInfo
+                                         usingAudioEmbedder:audioEmbedder
+                              expectedEmbeddingResultsCount:
+                                  kExpectedEmbeddingResultsCountForSpeechFiles];
   [MPPAudioEmbedderTests
-      assertResultsOfEmbedAudioClipWithFileInfo:kSpeech48KHzMonoFileInfo
-                             usingAudioEmbedder:audioEmbedder
-                                    isQuantized:options.quantize
-                  expectedEmbeddingResultsCount:kExpectedEmbeddingResultsCountForSpeechFiles];
+      assertResultsOfNonQuantizedEmbedAudioClipWithFileInfo:kSpeech48KHzMonoFileInfo
+                                         usingAudioEmbedder:audioEmbedder
+                              expectedEmbeddingResultsCount:
+                                  kExpectedEmbeddingResultsCountForSpeechFiles];
 }
 
 - (void)testEmbedWithSilenceSucceeds {
@@ -180,11 +180,10 @@ - (void)testEmbedWithSilenceSucceeds {
   MPPAudioEmbedderResult *result = [audioEmbedder embedAudioClip:audioData error:nil];
   XCTAssertNotNil(result);
 
-  const NSInteger expectedClassificationResultsCount = 1;
-
+  const NSInteger expectedEmbedderResultsCount = 1;
   [MPPAudioEmbedderTests assertAudioEmbedderResult:result
                                        isQuantized:options.quantize
-                  hasExpectedEmbeddingResultsCount:expectedClassificationResultsCount
+                     expectedEmbeddingResultsCount:expectedEmbedderResultsCount
                            expectedEmbeddingLength:kExpectedEmbeddingLength];
 }
 
@@ -225,6 +224,7 @@ - (void)testEmbedAfterCloseFailsInAudioClipsMode {
 }
 
 #pragma mark Running mode tests
+
 - (void)testCreateAudioEmbedderFailsWithDelegateInAudioClipsMode {
   MPPAudioEmbedderOptions *options =
       [MPPAudioEmbedderTests audioEmbedderOptionsWithModelFileInfo:kYamnetModelFileInfo];
@@ -337,26 +337,13 @@ - (void)testCreateAudioRecordWithInvalidChannelCountFails {
   AssertEqualErrors(error, expectedError);
 }
 
-#pragma mark MPPAudioEmbedderrStreamDelegate
+#pragma mark MPPAudioEmbedderStreamDelegate
 
 - (void)audioEmbedder:(MPPAudioEmbedder *)audioEmbedder
     didFinishEmbeddingWithResult:(MPPAudioEmbedderResult *)result
          timestampInMilliseconds:(NSInteger)timestampInMilliseconds
                            error:(NSError *)error {
-  // TODO: Test for results when stream mode inference tests are added.
-}
-
-#pragma mark Audio Data Initializers
-
-+ (NSArray<MPPTimestampedAudioData *> *)streamedAudioDataListforYamnet {
-  NSArray<MPPTimestampedAudioData *> *streamedAudioDataList =
-      [AVAudioFile streamedAudioBlocksFromAudioFileWithInfo:kSpeech16KHzMonoFileInfo
-                                           modelSampleCount:kYamnetSampleCount
-                                            modelSampleRate:kYamnetSampleRate];
-
-  XCTAssertEqual(streamedAudioDataList.count, 5);
-
-  return streamedAudioDataList;
+  // TODO: Add assertion for the result when stream mode inference tests are added.
 }
 
 #pragma mark Audio Embedder Initializers
@@ -399,6 +386,16 @@ + (void)assertCreateAudioEmbedderWithOptions:(MPPAudioEmbedderOptions *)options
 
 #pragma mark Results
 
++ (void)assertResultsOfNonQuantizedEmbedAudioClipWithFileInfo:(MPPFileInfo *)fileInfo
+                                           usingAudioEmbedder:(MPPAudioEmbedder *)audioEmbedder
+                                expectedEmbeddingResultsCount:
+                                    (NSInteger)expectedEmbeddingResultsCount {
+  [MPPAudioEmbedderTests assertResultsOfEmbedAudioClipWithFileInfo:fileInfo
+                                                usingAudioEmbedder:audioEmbedder
+                                                       isQuantized:NO
+                                     expectedEmbeddingResultsCount:expectedEmbeddingResultsCount];
+}
+
 + (void)assertResultsOfEmbedAudioClipWithFileInfo:(MPPFileInfo *)fileInfo
                                usingAudioEmbedder:(MPPAudioEmbedder *)audioEmbedder
                                       isQuantized:(BOOL)isQuantized
@@ -408,14 +405,14 @@ + (void)assertResultsOfEmbedAudioClipWithFileInfo:(MPPFileInfo *)fileInfo
 
   [MPPAudioEmbedderTests assertAudioEmbedderResult:result
                                        isQuantized:isQuantized
-                  hasExpectedEmbeddingResultsCount:expectedEmbeddingResultsCount
+                     expectedEmbeddingResultsCount:expectedEmbeddingResultsCount
                            expectedEmbeddingLength:kExpectedEmbeddingLength];
 }
 
 + (void)assertAudioEmbedderResult:(MPPAudioEmbedderResult *)result
-                         isQuantized:(BOOL)isQuantized
-    hasExpectedEmbeddingResultsCount:(NSInteger)expectedEmbeddingResultsCount
-             expectedEmbeddingLength:(NSInteger)expectedEmbeddingLength {
+                      isQuantized:(BOOL)isQuantized
+    expectedEmbeddingResultsCount:(NSInteger)expectedEmbeddingResultsCount
+          expectedEmbeddingLength:(NSInteger)expectedEmbeddingLength {
   XCTAssertEqual(result.embeddingResults.count, expectedEmbeddingResultsCount);
   for (MPPEmbeddingResult *embeddingResult in result.embeddingResults) {
     AssertEmbeddingResultHasOneEmbedding(embeddingResult);