annotations/Azaion.Annotator/Extensions/VLCFrameExtractor.cs

using System.Collections.Concurrent;
using System.IO;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using Azaion.Common.DTO.Config;
using LibVLCSharp.Shared;
using SkiaSharp;

namespace Azaion.Annotator.Extensions;

public class VLCFrameExtractor(LibVLC libVLC, AIRecognitionConfig config)
{
    private const uint RGBA_BYTES = 4;
    private const int PLAYBACK_RATE = 4;

    private uint _pitch; // Number of bytes per "line", aligned to x32.
    private uint _lines; // Number of lines in the buffer, aligned to x32.
    private uint _width; // Thumbnail width
    private uint _height; // Thumbnail height

    private MediaPlayer _mediaPlayer = null!;

    private TimeSpan _lastFrameTimestamp;
    private long _lastFrame;

    private static uint Align32(uint size)
    {
        if (size % 32 == 0)
            return size;
        return (size / 32 + 1) * 32;// Align on the next multiple of 32
    }

    private static SKBitmap? _currentBitmap;
    private static readonly ConcurrentQueue<FrameInfo> FramesQueue = new();
    private static long _frameCounter;

    public async IAsyncEnumerable<(TimeSpan Time, Stream Stream)> ExtractFrames(string mediaPath,
        [EnumeratorCancellation] CancellationToken manualCancellationToken = default)
    {
        var videoFinishedCancellationSource = new CancellationTokenSource();

        _mediaPlayer = new MediaPlayer(libVLC);
        _mediaPlayer.Stopped += (s, e) => videoFinishedCancellationSource.CancelAfter(1);

        using var media = new Media(libVLC, mediaPath);
        await media.Parse(cancellationToken: videoFinishedCancellationSource.Token);
        var videoTrack = media.Tracks.FirstOrDefault(x => x.Data.Video.Width != 0);
        _width = videoTrack.Data.Video.Width;
        _height = videoTrack.Data.Video.Height;

        _pitch = Align32(_width * RGBA_BYTES);
        _lines = Align32(_height);
        _mediaPlayer.SetRate(PLAYBACK_RATE);

        media.AddOption(":no-audio");
        _mediaPlayer.SetVideoFormat("RV32", _width, _height, _pitch);
        _mediaPlayer.SetVideoCallbacks(Lock, null, Display);

        _mediaPlayer.Play(media);
        _frameCounter = 0;
        var surface = SKSurface.Create(new SKImageInfo((int) _width, (int) _height));
        var videoFinishedCT = videoFinishedCancellationSource.Token;

        while ( !(FramesQueue.IsEmpty && videoFinishedCT.IsCancellationRequested || manualCancellationToken.IsCancellationRequested))
        {
            if (FramesQueue.TryDequeue(out var frameInfo))
            {
                if (frameInfo.Bitmap == null)
                    continue;

                surface.Canvas.DrawBitmap(frameInfo.Bitmap, 0, 0); // Effectively crops the original bitmap to get only the visible area

                using var outputImage = surface.Snapshot();
                using var data = outputImage.Encode(SKEncodedImageFormat.Jpeg, 85);
                using var ms = new MemoryStream();
                data.SaveTo(ms);

                yield return (frameInfo.Time, ms);

                frameInfo.Bitmap?.Dispose();
            }
            else
            {
                await Task.Delay(TimeSpan.FromSeconds(1), videoFinishedCT);
            }
        }
        FramesQueue.Clear(); //clear queue in case of manual stop
        _mediaPlayer.Stop();
        _mediaPlayer.Dispose();
    }

    private IntPtr Lock(IntPtr opaque, IntPtr planes)
    {
        _currentBitmap = new SKBitmap(new SKImageInfo((int)(_pitch / RGBA_BYTES), (int)_lines, SKColorType.Bgra8888));
        Marshal.WriteIntPtr(planes, _currentBitmap.GetPixels());
        return IntPtr.Zero;
    }

    private void Display(IntPtr opaque, IntPtr picture)
    {
        var playerTime = TimeSpan.FromMilliseconds(_mediaPlayer.Time);
        if (_lastFrameTimestamp != playerTime)
        {
            _lastFrame = _frameCounter;
            _lastFrameTimestamp = playerTime;
        }

        if (_frameCounter > 20 && _frameCounter % config.FramePeriodRecognition == 0)
        {
            var msToAdd = (_frameCounter - _lastFrame) * (_lastFrame == 0 ? 0 : _lastFrameTimestamp.TotalMilliseconds / _lastFrame);
            var time = _lastFrameTimestamp.Add(TimeSpan.FromMilliseconds(msToAdd));

            FramesQueue.Enqueue(new FrameInfo(time, _currentBitmap));
        }
        else
        {
            _currentBitmap?.Dispose();
        }

        _currentBitmap = null;
        _frameCounter++;
    }
}

public class FrameInfo(TimeSpan time, SKBitmap? bitmap)
{
    public TimeSpan Time { get; set; } = time;
    public SKBitmap? Bitmap { get; set; } = bitmap;
}