Internal change

PiperOrigin-RevId: 521909998

mediapipe/objc/BUILD

@@ -125,8 +125,10 @@ objc_library(
     visibility = ["//visibility:public"],
     deps = [
         "//third_party/apple_frameworks:AVFoundation",
+        "//third_party/apple_frameworks:CoreAudio",
         "//third_party/apple_frameworks:CoreVideo",
         "//third_party/apple_frameworks:Foundation",
+        "//third_party/apple_frameworks:MediaToolbox",
     ],
 )
 

mediapipe/objc/MPPInputSource.h

@@ -13,8 +13,11 @@
 // limitations under the License.
 
 #import <AVFoundation/AVFoundation.h>
+#import <CoreAudio/CoreAudioTypes.h>
 #import <Foundation/Foundation.h>
 
+NS_ASSUME_NONNULL_BEGIN
+
 @class MPPInputSource;
 
 /// A delegate that can receive frames from a source.
@@ -31,7 +34,7 @@
                 timestamp:(CMTime)timestamp
                fromSource:(MPPInputSource*)source;
 
-// Provides the delegate with a new depth frame data
+// Provides the delegate with new depth frame data.
 @optional
 - (void)processDepthData:(AVDepthData*)depthData
                timestamp:(CMTime)timestamp
@@ -40,6 +43,23 @@
 @optional
 - (void)videoDidPlayToEnd:(CMTime)timestamp;
 
+// Provides the delegate with the format of the audio track to be played.
+@optional
+- (void)willStartPlayingAudioWithFormat:(const AudioStreamBasicDescription*)format
+                             fromSource:(MPPInputSource*)source;
+
+// Lets the delegate know that there is no audio track despite audio playback
+// having been requested (or that audio playback failed for other reasons).
+@optional
+- (void)noAudioAvailableFromSource:(MPPInputSource*)source;
+
+// Provides the delegate with a new audio packet.
+@optional
+- (void)processAudioPacket:(const AudioBufferList*)audioPacket
+                 numFrames:(CMItemCount)numFrames
+                 timestamp:(CMTime)timestamp
+                fromSource:(MPPInputSource*)source;
+
 @end
 
 /// Abstract class for a video source.
@@ -68,3 +88,5 @@
 - (void)stop;
 
 @end
+
+NS_ASSUME_NONNULL_END

mediapipe/objc/MPPPlayerInputSource.h

@@ -18,7 +18,15 @@
 /// Not meant for batch processing of video.
 @interface MPPPlayerInputSource : MPPInputSource
 
-/// Designated initializer.
+/// Initializes the video source with optional audio processing.
+///
+/// @param video The video asset to play.
+/// @param audioProcessingEnabled If set, indicates that the (first) audio track
+///        should be processed if it exists, and the corresponding methods for
+///        audio will be invoked on the @c delegate.
+- (instancetype)initWithAVAsset:(AVAsset*)video audioProcessingEnabled:(BOOL)audioProcessingEnabled;
+
+/// Initializes the video source to process @c video with audio processing disabled.
 - (instancetype)initWithAVAsset:(AVAsset*)video;
 
 /// Skip into video @c time from beginning (time 0), within error of +/- tolerance to closest time.

mediapipe/objc/MPPPlayerInputSource.m

@@ -13,11 +13,13 @@
 // limitations under the License.
 
 #import <CoreVideo/CoreVideo.h>
+#import <MediaToolbox/MediaToolbox.h>
 
 #import "MPPPlayerInputSource.h"
 #if !TARGET_OS_OSX
 #import "mediapipe/objc/MPPDisplayLinkWeakTarget.h"
 #endif
+#import "mediapipe/objc/MPPInputSource.h"
 
 @implementation MPPPlayerInputSource {
   AVAsset* _video;
@@ -35,7 +37,53 @@
   BOOL _playing;
 }
 
+void InitAudio(MTAudioProcessingTapRef tap, void* clientInfo, void** tapStorageOut) {
+  // `clientInfo` comes as a user-defined argument through
+  // `MTAudioProcessingTapCallbacks`; we pass our `MPPPlayerInputSource`
+  // there. Tap processing functions allow for user-defined "storage" - we just
+  // treat our input source as such.
+  *tapStorageOut = clientInfo;
+}
+
+void PrepareAudio(MTAudioProcessingTapRef tap, CMItemCount maxFrames,
+                  const AudioStreamBasicDescription* audioFormat) {
+  // See `InitAudio`.
+  MPPPlayerInputSource* source =
+      (__bridge MPPPlayerInputSource*)MTAudioProcessingTapGetStorage(tap);
+  if ([source.delegate respondsToSelector:@selector(willStartPlayingAudioWithFormat:fromSource:)]) {
+    [source.delegate willStartPlayingAudioWithFormat:audioFormat fromSource:source];
+  }
+}
+
+void ProcessAudio(MTAudioProcessingTapRef tap, CMItemCount numberFrames,
+                  MTAudioProcessingTapFlags flags, AudioBufferList* bufferListInOut,
+                  CMItemCount* numberFramesOut, MTAudioProcessingTapFlags* flagsOut) {
+  CMTimeRange timeRange;
+  OSStatus status = MTAudioProcessingTapGetSourceAudio(tap, numberFrames, bufferListInOut, flagsOut,
+                                                       &timeRange, numberFramesOut);
+  if (status != 0) {
+    NSLog(@"Error from GetSourceAudio: %ld", (long)status);
+    return;
+  }
+
+  // See `InitAudio`.
+  MPPPlayerInputSource* source =
+      (__bridge MPPPlayerInputSource*)MTAudioProcessingTapGetStorage(tap);
+  if ([source.delegate respondsToSelector:@selector(processAudioPacket:
+                                                             numFrames:timestamp:fromSource:)]) {
+    [source.delegate processAudioPacket:bufferListInOut
+                              numFrames:numberFrames
+                              timestamp:timeRange.start
+                             fromSource:source];
+  }
+}
+
 - (instancetype)initWithAVAsset:(AVAsset*)video {
+  return [self initWithAVAsset:video audioProcessingEnabled:NO];
+}
+
+- (instancetype)initWithAVAsset:(AVAsset*)video
+         audioProcessingEnabled:(BOOL)audioProcessingEnabled {
   self = [super init];
   if (self) {
     _video = video;
@@ -67,6 +115,11 @@
     CVDisplayLinkStop(_videoDisplayLink);
     CVDisplayLinkSetOutputCallback(_videoDisplayLink, renderCallback, (__bridge void*)self);
 #endif  // TARGET_OS_OSX
+
+    if (audioProcessingEnabled) {
+      [self setupAudioPlayback];
+    }
+
     _videoPlayer = [AVPlayer playerWithPlayerItem:_videoItem];
     _videoPlayer.actionAtItemEnd = AVPlayerActionAtItemEndNone;
     NSNotificationCenter* center = [NSNotificationCenter defaultCenter];
@@ -88,6 +141,47 @@
   return self;
 }
 
+- (void)setupAudioPlayback {
+  bool have_audio = false;
+  NSArray<AVAssetTrack*>* audioTracks =
+      [_video tracksWithMediaCharacteristic:AVMediaCharacteristicAudible];
+  if (audioTracks.count != 0) {
+    // We always limit ourselves to the first audio track if there are
+    // multiple (which is a rarity) - note that it can still be e.g. stereo.
+    AVAssetTrack* audioTrack = audioTracks[0];
+    MTAudioProcessingTapCallbacks audioCallbacks;
+    audioCallbacks.version = kMTAudioProcessingTapCallbacksVersion_0;
+    audioCallbacks.clientInfo = (__bridge void*)(self);
+    audioCallbacks.init = InitAudio;
+    audioCallbacks.prepare = PrepareAudio;
+    audioCallbacks.process = ProcessAudio;
+    audioCallbacks.unprepare = NULL;
+    audioCallbacks.finalize = NULL;
+
+    MTAudioProcessingTapRef audioTap;
+    OSStatus status =
+        MTAudioProcessingTapCreate(kCFAllocatorDefault, &audioCallbacks,
+                                   kMTAudioProcessingTapCreationFlag_PreEffects, &audioTap);
+    if (status == noErr && audioTap != NULL) {
+      AVMutableAudioMixInputParameters* audioMixInputParams =
+          [AVMutableAudioMixInputParameters audioMixInputParametersWithTrack:audioTrack];
+      audioMixInputParams.audioTapProcessor = audioTap;
+      CFRelease(audioTap);
+
+      AVMutableAudioMix* audioMix = [AVMutableAudioMix audioMix];
+
+      audioMix.inputParameters = @[ audioMixInputParams ];
+      _videoItem.audioMix = audioMix;
+      have_audio = true;
+    } else {
+      NSLog(@"Error %ld when trying to create the audio processing tap", (long)status);
+    }
+  }
+  if (!have_audio && [self.delegate respondsToSelector:@selector(noAudioAvailableFromSource:)]) {
+    [self.delegate noAudioAvailableFromSource:self];
+  }
+}
+
 - (void)start {
   [_videoPlayer play];
   _playing = YES;