Make Calculator more thread-safe (#25)

* YEP implement @LeNitrous 's suggestion * lock * Fix stream handle nullable issue * Holistic calculators Co-authored-by: Speykious <speykious@gmail.com>
cosyneco · Mar 4, 2022 · 9dcb9c3 · 9dcb9c3
1 parent cc7ed1e
commit 9dcb9c3
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 9 deletions.
diff --git a/Mediapipe.Net.Examples.FaceMeshGpu/Program.cs b/Mediapipe.Net.Examples.FaceMeshGpu/Program.cs
@@ -82,7 +82,6 @@ public static void Main(string[] args)
                     cFrame.Width, cFrame.Height, cFrame.WidthStep, cFrame.RawData);
 
                 using ImageFrame img = calculator.Send(imgframe);
-                imgframe.Dispose();
             }
         }
 

diff --git a/Mediapipe.Net.Examples.OsuFrameworkVisualTests/MediapipeDrawable.cs b/Mediapipe.Net.Examples.OsuFrameworkVisualTests/MediapipeDrawable.cs
@@ -64,7 +64,6 @@ protected override unsafe void Update()
             ImageFrame imgFrame = new ImageFrame(ImageFormat.Srgba,
                 cFrame.Width, cFrame.Height, cFrame.WidthStep, cFrame.RawData);
             using ImageFrame outImgFrame = calculator.Send(imgFrame);
-            imgFrame.Dispose();
 
             var span = new ReadOnlySpan<byte>(outImgFrame.MutablePixelData, outImgFrame.Height * outImgFrame.WidthStep);
             var pixelData = Image.LoadPixelData<Rgba32>(span, cFrame.Width, cFrame.Height);

diff --git a/Mediapipe.Net/Calculators/Calculator.cs b/Mediapipe.Net/Calculators/Calculator.cs
@@ -29,7 +29,7 @@ public abstract class Calculator<TPacket, T> : Disposable
         protected readonly string? SecondaryOutputStream;
 
         protected readonly CalculatorGraph Graph;
-        private GCHandle observeStreamHandle;
+        private GCHandle? observeStreamHandle;
 
         /// <summary>
         /// Triggered every time the calculator returns a secondary output.
@@ -53,7 +53,8 @@ protected Calculator(string graphPath, string? secondaryOutputStream = null)
                     T secondaryOutput = packet.Get();
                     OnResult?.Invoke(this, secondaryOutput);
                     return Status.Ok();
-                }, out observeStreamHandle).AssertOk();
+                }, out GCHandle handle).AssertOk();
+                observeStreamHandle = handle;
             }
         }
 
@@ -70,12 +71,18 @@ protected Calculator(string graphPath, string? secondaryOutputStream = null)
         /// </summary>
         /// <remarks>If the input <see cref="ImageFrame"/> doesn't get disposed after being sent, MediaPipe will crash.</remarks>
         /// <param name="frame">The frame that MediaPipe should process.</param>
+        /// <param name="disposeSourceFrame">Whether or not to dispose the source frame.</param>
         /// <returns>An <see cref="ImageFrame"/> with the contents of the source <see cref="ImageFrame"/> and the MediaPipe solution drawn.</returns>
-        public ImageFrame Send(ImageFrame frame)
+        public ImageFrame Send(ImageFrame frame, bool disposeSourceFrame = true)
         {
-            ImageFrame outFrame = SendFrame(frame);
-            CurrentFrame++;
-            return outFrame;
+            lock (frame)
+            {
+                ImageFrame outFrame = SendFrame(frame);
+                CurrentFrame++;
+                if (disposeSourceFrame)
+                    frame.Dispose();
+                return outFrame;
+            }
         }
 
         /// <summary>
@@ -89,7 +96,7 @@ protected override void DisposeManaged()
             Graph.WaitUntilDone();
             Graph.Dispose();
 
-            observeStreamHandle.Free();
+            observeStreamHandle?.Free();
         }
     }
 }
diff --git a/Mediapipe.Net/Calculators/HolisticCpuCalculator.cs b/Mediapipe.Net/Calculators/HolisticCpuCalculator.cs
@@ -0,0 +1,16 @@
+// Copyright (c) homuler and The Vignette Authors
+// This file is part of MediaPipe.NET.
+// MediaPipe.NET is licensed under the MIT License. See LICENSE for details.
+
+using Mediapipe.Net.Framework.Packet;
+
+namespace Mediapipe.Net.Calculators
+{
+    public sealed class HolisticCpuCalculator : CpuCalculator<BoolPacket, bool>
+    {
+        public HolisticCpuCalculator() : base(
+            graphPath: "mediapipe/graphs/holistic_tracking/holistic_tracking_cpu.pbtxt")
+        {
+        }
+    }
+}
diff --git a/Mediapipe.Net/Calculators/HolisticGpuCalculator.cs b/Mediapipe.Net/Calculators/HolisticGpuCalculator.cs
@@ -0,0 +1,16 @@
+// Copyright (c) homuler and The Vignette Authors
+// This file is part of MediaPipe.NET.
+// MediaPipe.NET is licensed under the MIT License. See LICENSE for details.
+
+using Mediapipe.Net.Framework.Packet;
+
+namespace Mediapipe.Net.Calculators
+{
+    public sealed class HolisticGpuCalculator : CpuCalculator<BoolPacket, bool>
+    {
+        public HolisticGpuCalculator() : base(
+            graphPath: "mediapipe/graphs/holistic_tracking/holistic_tracking_gpu.pbtxt")
+        {
+        }
+    }
+}