WPF-MediaKit video screen grabbing helper.

VideoScreenGrabber.cs

using System;
using System.IO;
using System.Threading;
using System.Threading.Tasks;
using System.Windows.Interop;
using WPFMediaKit.DirectShow.Controls;
using WPFMediaKit.DirectShow.MediaPlayers;

namespace Test_Application
{
    /// <summary>
    /// Class to simplify video screen grabbing by the async tasks.
    /// </summary>
    public class VideoScreenGrabber : IDisposable
    {
        private TaskCompletionSource<bool> taskCompletionOpen;
        private TaskCompletionSource<IntPtr> taskCompletionGrab;
        private CancellationTokenRegistration cancellationRegistrationGrab;

        public MediaUriPlayer Player { get; private set; }

        /// <summary>
        /// Current back buffer.
        /// </summary>
        public IntPtr BackBuffer { get; private set; }

        /// <summary>
        /// Media duration in 100ns units.
        /// </summary>
        public long MediaDuration
        {
            get
            {
                CheckPlayer();
                return Player.Duration;
            }
        }

        /// <summary>
        /// Media duration [sec].
        /// </summary>
        public double MediaDurationSecond
            => (double)MediaDuration / MediaPlayerBase.DSHOW_ONE_SECOND_UNIT;

        /// <summary>
        /// Busy with grabbing.
        /// </summary>
        public bool IsGrabbing
            => taskCompletionGrab != null;

        /// <summary>
        /// Open given source. Fails with exception if the file cannot be played.
        /// </summary>
        public Task Open(Uri source)
        {
            if (Player != null)
                throw new ArgumentException("Cannot open twice!");
            taskCompletionOpen = new TaskCompletionSource<bool>();

            Player = new MediaUriPlayer();
            Player.EnsureThread(ApartmentState.MTA);
            Player.MediaOpened += MediaUriPlayer_MediaOpened;
            Player.MediaFailed += MediaUriPlayer_MediaFailed;
            Player.NewAllocatorFrame += Player_NewAllocatorFrame;
            Player.NewAllocatorSurface += Player_NewAllocatorSurface;

            Player.Dispatcher.BeginInvoke(new Action(() =>
            {
                Player.AudioDecoder = null;
                Player.AudioRenderer = null;
                Player.Source = source;
            }));

            return taskCompletionOpen.Task;
        }

        /// <summary>
        /// Grab a buffer at given position [sec].
        /// See <see cref="GrabAtPosition(long)"/>.
        /// </summary>
        public Task<IntPtr> GrabAtSecond(double second)
            => GrabAtPosition((long)(second * MediaPlayerBase.DSHOW_ONE_SECOND_UNIT));

        /// <summary>
        /// Grab a buffer at given position.
        /// </summary>
        /// <param name="position">Video position in 100ns units.</param>
        /// <param name="cancellationToken">CancellationToken, may be used for timeout.</param>
        /// <returns>Buffer for D3DImage.</returns>
        public Task<IntPtr> GrabAtPosition(long position, CancellationToken cancellationToken = default(CancellationToken))
        {
            CheckPlayer();
            if (taskCompletionGrab != null)
                throw new InvalidOperationException("Still grabbing previous frame.");
            if (position < 0)
                throw new ArgumentException("position negative.");
            if (position > MediaDuration)
                throw new ArgumentException("position beyond the media duration.");

            taskCompletionGrab = new TaskCompletionSource<IntPtr>();
            cancellationRegistrationGrab = cancellationToken.Register(CancelGrab);
            Player.Dispatcher.BeginInvoke(new Action(() =>
            {
                Player.MediaPosition = position;
                Player.Pause();
            }));

            return taskCompletionGrab.Task;
        }

        public void CancelGrab()
        {
            taskCompletionGrab?.TrySetCanceled();
            taskCompletionGrab = null;
            cancellationRegistrationGrab.Dispose();
        }

        private void CheckPlayer()
        {
            if (Player == null)
                throw new InvalidOperationException("Player not opened, call Open() first.");
        }

        private void MediaUriPlayer_MediaFailed(object sender, MediaFailedEventArgs e)
        {
            Exception exc = e.Exception;
            if (exc == null)
                exc = new WPFMediaKit.WPFMediaKitException(e.Message);
            taskCompletionOpen.TrySetException(exc);
        }

        private void MediaUriPlayer_MediaOpened()
            => taskCompletionOpen.TrySetResult(true);

        private void Player_NewAllocatorSurface(object sender, IntPtr pSurface)
            => BackBuffer = pSurface;

        private void Player_NewAllocatorFrame()
        {
            if (taskCompletionGrab == null)
                return;
            taskCompletionGrab.TrySetResult(BackBuffer);
            // not exactly thread safe, but do the job in common scenarios
            taskCompletionGrab = null;
        }

        public void Dispose()
        {
            Dispose(true);
        }

        protected virtual void Dispose(bool disposing)
        {
            if (!disposing)
                return;
            CancelGrab();
            if (Player != null)
            {
                Player.Dispose();
                Player = null;
            }
        }
    }

    /// <summary>
    /// Helper class for one time frame grabbing.
    /// </summary>
    public static class VideoScreenGrabberUtils
    {
        private const int TIMEOUT_MS = 5000;

        /// <summary>
        /// Grab a D3DImage with timeout.
        /// </summary>
        /// <param name="source">Video source Uri.</param>
        /// <param name="position">Position in DSHOW_ONE_SECOND_UNIT units.</param>
        /// <param name="timeout">Timeout in millis. If zero, then wait infinitely.</param>
        /// <returns>Created D3DImage.</returns>
        /// <exception cref="TimeoutException">When the opearion has timeed out.</exception>
        public static async Task<D3DImage> GrabScreenAtPosition(Uri source, long position, int timeout = TIMEOUT_MS)
        {
            using (VideoScreenGrabber grabber = new VideoScreenGrabber())
            {
                await grabber.Open(source);
                long duration = grabber.MediaDuration;
                // if the video is too short, grab the screen at half
                if (position > duration - (MediaPlayerBase.DSHOW_ONE_SECOND_UNIT / 2))
                    position = duration / 2;
                IntPtr buffer;
                if (timeout > 0)
                {
                    try
                    {
                        buffer = await grabber.GrabAtPosition(position, new CancellationTokenSource(timeout).Token);
                    }
                    catch (TaskCanceledException)
                    {
                        throw new TimeoutException("The operation has timed-out.");
                    }
                }
                else
                {
                    buffer = await grabber.GrabAtPosition(position);
                }
                D3DImage d3d = new D3DImage();
                D3DImageUtils.SetBackBufferWithLock(d3d, buffer);
                return d3d;
            }
        }
    }
}