fix: dispose native resources of segmentation masks properly (#1166)

* fix: avoid SIGSEGV when reading the internal image * fix: dispose native resources of segmentation masks * fix: avoid data race for PoseLandmarker
homuler · Mar 14, 2024 · 019a655 · 019a655
1 parent 9e2d601
commit 019a655
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 10 deletions.
diff --git a/Packages/com.github.homuler.mediapipe/Runtime/Scripts/Framework/PacketGetterExtension.cs b/Packages/com.github.homuler.mediapipe/Runtime/Scripts/Framework/PacketGetterExtension.cs
@@ -308,7 +308,7 @@ public static void Get(this Packet<List<Image>> packet, List<Image> value)
 
       foreach (var imagePtr in imageArray.AsReadOnlySpan())
       {
-        value.Add(new Image(imagePtr, false));
+        value.Add(new Image(imagePtr, true));
       }
     }
 

diff --git a/...er.mediapipe/Runtime/Scripts/Unity/Annotation/PoseLandmarkerResultAnnotationController.cs b/...er.mediapipe/Runtime/Scripts/Unity/Annotation/PoseLandmarkerResultAnnotationController.cs
@@ -14,6 +14,7 @@ public class PoseLandmarkerResultAnnotationController : AnnotationController<Mul
   {
     [SerializeField] private bool _visualizeZ = false;
 
+    private readonly object _currentTargetLock = new object();
     private PoseLandmarkerResult _currentTarget;
 
     public void InitScreen(int maskWidth, int maskHeight) => annotation.InitMask(maskWidth, maskHeight);
@@ -30,7 +31,7 @@ public void DrawNow(PoseLandmarkerResult target)
 
     protected void UpdateCurrentTarget(PoseLandmarkerResult newTarget)
     {
-      if (IsTargetChanged(newTarget, _currentTarget))
+      lock (_currentTargetLock)
       {
         newTarget.CloneTo(ref _currentTarget);
         isStale = true;
@@ -39,17 +40,20 @@ protected void UpdateCurrentTarget(PoseLandmarkerResult newTarget)
 
     protected override void SyncNow()
     {
-      isStale = false;
-      if (_currentTarget.segmentationMasks != null)
+      lock (_currentTargetLock)
       {
-        ReadMask(_currentTarget.segmentationMasks);
-        // TODO: stop disposing masks here
-        foreach (var mask in _currentTarget.segmentationMasks)
+        isStale = false;
+        if (_currentTarget.segmentationMasks != null)
         {
-          mask.Dispose();
+          ReadMask(_currentTarget.segmentationMasks);
+          // TODO: stop disposing masks here
+          foreach (var mask in _currentTarget.segmentationMasks)
+          {
+            mask.Dispose();
+          }
         }
+        annotation.Draw(_currentTarget.poseLandmarks, _visualizeZ);
       }
-      annotation.Draw(_currentTarget.poseLandmarks, _visualizeZ);
     }
   }
 }
diff --git a/mediapipe_api/framework/formats/image.cc b/mediapipe_api/framework/formats/image.cc
@@ -133,7 +133,7 @@ MpReturnCode mp_Packet__GetImageVector(mediapipe::Packet* packet, mp_api::Struct
     auto data = new mediapipe::Image*[size];
 
     for (auto i = 0; i < size; ++i) {
-      data[i] = new mediapipe::Image(std::move(vec[i]));
+      data[i] = new mediapipe::Image(vec[i].GetImageFrameSharedPtr());
     }
     value_out->data = data;
     value_out->size = static_cast<int>(size);