BlazePoseOfflineCpuCalculator.cs

// Inspired by CPU Calculator from ImagePipe.NET, but with the secondary output stream polled rather than event-based.

#nullable enable
using System;
using System.Diagnostics;
using System.IO;
using System.Runtime.InteropServices;
using FluentAssertions;
using Mediapipe.Net.Core;
using Mediapipe.Net.Framework;
using Mediapipe.Net.Framework.Format;
using Mediapipe.Net.Framework.Packet;
using Mediapipe.Net.Framework.Protobuf;
using Microsoft.Extensions.Logging;

namespace myapp.Services.MediaPipe
{
	/// <summary>
	/// Based on MediaPipe.Net calculators, but using pollers rather than notification, so we can pass one packet in at a time.
	/// </summary>
	public class BlazePoseOfflineCpuCalculator : Disposable
	{
	#region Constants
		private string _graphPath = "mediapipe/graphs/pose_tracking/vf_pose_tracking_cpu.pbtxt";
		private string _inputStreamName = "input_video";
		private string _primaryOutputName = "output_video";
		private string _secondaryOutputName = "pose_landmarks";
	#endregion

	#region Fields
		private readonly ILogger<BlazePoseOfflineCpuCalculator> _logger;
		private CalculatorGraph? _graph;
		private OutputStreamPoller<ImageFrame>? _framePoller;
		private OutputStreamPoller<NormalizedLandmarkList>? _secondaryPoller;
	#endregion

	#region Properties
		/// <summary>
		/// The number of the current processed frame.
		/// </summary>
		public long CurrentTimestamp { get; private set; } = 0;
	#endregion


	#region Lifetime
		public BlazePoseOfflineCpuCalculator(ILogger<BlazePoseOfflineCpuCalculator> logger)
		{
			_logger = logger;
		}

		// Add an init step, because there's too much that can go wrong in the constructor otherwise
		// We'll do it automatically the first time run() is called
		private void Initialize()
		{
			var graphText = File.ReadAllText(_graphPath);
			_logger.LogDebug("Graph text: {Graph}", graphText);
			_graph = new CalculatorGraph(graphText);

			// Try setting up pollers for both output streams 
			_framePoller = _graph.AddOutputStreamPoller<ImageFrame>(_primaryOutputName).Value();
			_secondaryPoller = _graph.AddOutputStreamPoller<NormalizedLandmarkList>(_secondaryOutputName).Value();
		}

		protected override void DisposeManaged()
		{
			base.DisposeManaged();
			_framePoller?.Dispose();
			_secondaryPoller?.Dispose();
		}
	#endregion

	#region Methods
		public void Run()
		{
			if (_graph == null)
			{
				Initialize();
			}

			if (_graph == null)
			{
				throw new NullReferenceException("Graph should be initialized");
			}
			
			_graph.StartRun().AssertOk();
		}

		public (ImageFrame? primaryOutput, NormalizedLandmarkList? secondaryOutput) SendFrame(ImageFrame frame)
		{
			ImageFrame? primaryOutput = null;
			NormalizedLandmarkList? secondaryOutput = null;

			// I think because we're new'ing this with the current timestamp, it will use the bits in frame directly rather than copy them.
			using var packet = new ImageFramePacket(frame, new Timestamp(CurrentTimestamp++));

			var stopwatch = new Stopwatch();
			stopwatch.Start();
			_graph!.AddPacketToInputStream(_inputStreamName, packet).AssertOk();
			var timeToAddInputPacket = stopwatch.ElapsedMilliseconds;

			var outPacket = new ImageFramePacket();
			// Poll for the primary output
			if (_framePoller!.Next(outPacket))
			{
				_logger.LogDebug("framePoller returned true with packet timestamp: {Timestamp}", outPacket.Timestamp().Value);
				var outFrame = outPacket.Get();
				primaryOutput = outFrame;
			}
			else
			{
				// Dispose the unused outPacket
				_logger.LogDebug("framePoller returned false, disposing");
				outPacket.Dispose();
			}

			var secondaryPacket = new NormalizedLandmarkListPacket();
			// Poll for the secondary output
			if (_secondaryPoller!.Next(secondaryPacket))
			{
				var timeToPollForSecondary = stopwatch.ElapsedMilliseconds;

				_logger.LogDebug("secondaryPoller returned true with packet timestamp: {Timestamp}", secondaryPacket.Timestamp().Value);
				secondaryOutput = secondaryPacket.Get();
				var timeToGetSecondary = stopwatch.ElapsedMilliseconds;
				_logger.LogDebug("Timing: \r\n Time to add input packet: {T1}\r\n Time to poll for secondary: {T2} \r\r Time to get secondary: {T3}",
					timeToAddInputPacket, timeToPollForSecondary, timeToGetSecondary);
			}
			else
			{
				// Dispose the unused outPacket
				_logger.LogDebug("secondaryPoller returned false, disposing");
				secondaryPacket.Dispose();
			}

			stopwatch.Stop();
			return (primaryOutput, secondaryOutput);
		}
	#endregion
	}
}