From 9dcb9c34923ea7566c9643afa233a9ba8cc398ea Mon Sep 17 00:00:00 2001
From: "Adwyzz (OLED Edition)" <46694241+adryzz@users.noreply.github.com>
Date: Fri, 4 Mar 2022 15:34:25 +0000
Subject: [PATCH] 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>
---
 Mediapipe.Net.Examples.FaceMeshGpu/Program.cs |  1 -
 .../MediapipeDrawable.cs                      |  1 -
 Mediapipe.Net/Calculators/Calculator.cs       | 21 ++++++++++++-------
 .../Calculators/HolisticCpuCalculator.cs      | 16 ++++++++++++++
 .../Calculators/HolisticGpuCalculator.cs      | 16 ++++++++++++++
 5 files changed, 46 insertions(+), 9 deletions(-)
 create mode 100644 Mediapipe.Net/Calculators/HolisticCpuCalculator.cs
 create mode 100644 Mediapipe.Net/Calculators/HolisticGpuCalculator.cs

diff --git a/Mediapipe.Net.Examples.FaceMeshGpu/Program.cs b/Mediapipe.Net.Examples.FaceMeshGpu/Program.cs
index e9f7e14a..bf1c78a0 100644
--- 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
index 34bc2037..6caa8f65 100644
--- 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
index c94a9699..f7d75675 100644
--- 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
new file mode 100644
index 00000000..e482fe9b
--- /dev/null
+++ 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
new file mode 100644
index 00000000..b8c213be
--- /dev/null
+++ 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")
+        {
+        }
+    }
+}