Updated iOS image classifier and object detector callback invocations to avoid retain cycles

mediapipe/tasks/ios/vision/image_classifier/sources/MPPImageClassifier.mm

@@ -85,22 +85,34 @@ static NSString *const kTaskGraphName =
 
     if (options.imageClassifierDelegate) {
       _imageClassifierDelegate = options.imageClassifierDelegate;
+
+      // Capturing `self` as weak in order to avoid `self` being kept in memory
+      // and cause a retain cycle, after self is set to `nil`.
+      MPPImageClassifier *__weak weakSelf = self;
+      dispatch_queue_t callbackQueue =
+          dispatch_queue_create("com.mediapipe.tasks.imageClassifierCallbackQueue", NULL);
       packetsCallback = [=](absl::StatusOr<PacketMap> status_or_packets) {
-        NSError *callbackError = nil;
+        // Check to ensure that the delegate method is not called on a nil object
-        if (![MPPCommonUtils checkCppError:status_or_packets.status() toError:&callbackError]) {
+        // leading to a segmentation fault, we check `weakSelf` is `nil` before
-          if ([_imageClassifierDelegate
+        // performing any processing.
+        if (!weakSelf ||
+            ![weakSelf.imageClassifierDelegate
                 respondsToSelector:@selector
                 (imageClassifier:
                     didFinishClassificationWithResult:timestampInMilliseconds:error:)]) {
-            [_imageClassifierDelegate imageClassifier:self
-                    didFinishClassificationWithResult:nil
-                              timestampInMilliseconds:Timestamp::Unset().Value()
-                                                error:callbackError];
-          }
           return;
         }
 
-        PacketMap &outputPacketMap = status_or_packets.value();
+        NSError *callbackError = nil;
+        if (![MPPCommonUtils checkCppError:status_or_packets.status() toError:&callbackError]) {
+          [weakSelf.imageClassifierDelegate imageClassifier:weakSelf
+                          didFinishClassificationWithResult:nil
+                                    timestampInMilliseconds:Timestamp::Unset().Value()
+                                                      error:callbackError];
+          return;
+        }
+
+        PacketMap outputPacketMap = status_or_packets.value();
         if (outputPacketMap[kImageOutStreamName.cppString].IsEmpty()) {
           return;
         }
@@ -109,18 +121,14 @@ static NSString *const kTaskGraphName =
             [MPPImageClassifierResult imageClassifierResultWithClassificationsPacket:
                                           outputPacketMap[kClassificationsStreamName.cppString]];
 
-        if ([_imageClassifierDelegate
+        [weakSelf.imageClassifierDelegate
-                respondsToSelector:@selector
+                              imageClassifier:weakSelf
-                (imageClassifier:
-                    didFinishClassificationWithResult:timestampInMilliseconds:error:)]) {
-          [_imageClassifierDelegate imageClassifier:self
             didFinishClassificationWithResult:result
                       timestampInMilliseconds:outputPacketMap[kImageOutStreamName.cppString]
                                                   .Timestamp()
                                                   .Value() /
                                               kMicroSecondsPerMilliSecond
                                         error:callbackError];
-        }
       };
     }
 

mediapipe/tasks/ios/vision/object_detector/sources/MPPObjectDetector.mm

@@ -81,17 +81,23 @@ static NSString *const kTaskGraphName = @"mediapipe.tasks.vision.ObjectDetectorG
 
     if (options.objectDetectorDelegate) {
       _objectDetectorDelegate = options.objectDetectorDelegate;
+
+      // Capturing `self` as weak in order to avoid `self` being kept in memory
+      // and cause a retain cycle, after self is set to `nil`.
+      MPPObjectDetector *__weak weakSelf = self;
       packetsCallback = [=](absl::StatusOr<PacketMap> statusOrPackets) {
-        NSError *callbackError = nil;
+        if ([!weakSelf && weakSelf.objectDetectorDelegate
-        if (![MPPCommonUtils checkCppError:statusOrPackets.status() toError:&callbackError]) {
-          if ([_objectDetectorDelegate
                 respondsToSelector:@selector
                 (objectDetector:didFinishDetectionWithResult:timestampInMilliseconds:error:)]) {
-            [_objectDetectorDelegate objectDetector:self
+          return;
+        }
+
+        NSError *callbackError = nil;
+        if (![MPPCommonUtils checkCppError:statusOrPackets.status() toError:&callbackError]) {
+          [_objectDetectorDelegate objectDetector:weakSelf
                      didFinishDetectionWithResult:nil
                           timestampInMilliseconds:Timestamp::Unset().Value()
                                             error:callbackError];
-          }
           return;
         }
 
@@ -103,17 +109,15 @@ static NSString *const kTaskGraphName = @"mediapipe.tasks.vision.ObjectDetectorG
         MPPObjectDetectionResult *result = [MPPObjectDetectionResult
             objectDetectionResultWithDetectionsPacket:statusOrPackets.value()[kDetectionsStreamName
                                                                                   .cppString]];
-        if ([_objectDetectorDelegate
+
-                respondsToSelector:@selector
+        [weakSelf.objectDetectorDelegate
-                (objectDetector:didFinishDetectionWithResult:timestampInMilliseconds:error:)]) {
+                          objectDetector:weakSelf
-          [_objectDetectorDelegate objectDetector:self
             didFinishDetectionWithResult:result
                  timestampInMilliseconds:outputPacketMap[kImageOutStreamName.cppString]
                                              .Timestamp()
                                              .Value() /
                                          kMicroSecondsPerMilliSecond
                                    error:callbackError];
-        }
       };
     }