annotations/Azaion.Annotator/AnnotatorEventHandler.cs

using System.IO;
using System.Reflection.Metadata;
using System.Windows;
using System.Windows.Input;
using Azaion.Annotator.DTO;
using Azaion.Common;
using Azaion.Common.DTO;
using Azaion.Common.DTO.Config;
using Azaion.Common.DTO.Queue;
using Azaion.Common.Events;
using Azaion.Common.Extensions;
using Azaion.Common.Services;
using LibVLCSharp.Shared;
using MediatR;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using MediaPlayer = LibVLCSharp.Shared.MediaPlayer;

namespace Azaion.Annotator;

public class AnnotatorEventHandler(
    LibVLC libVLC,
    MediaPlayer mediaPlayer,
    Annotator mainWindow,
    FormState formState,
    AnnotationService annotationService,
    ILogger<AnnotatorEventHandler> logger,
    IOptions<DirectoriesConfig> dirConfig,
    IInferenceService inferenceService)
    :
            INotificationHandler<KeyEvent>,
            INotificationHandler<AnnClassSelectedEvent>,
            INotificationHandler<AnnotatorControlEvent>,
            INotificationHandler<VolumeChangedEvent>,
            INotificationHandler<AnnotationsDeletedEvent>
{
    private const int STEP = 20;
    private const int LARGE_STEP = 5000;
    private const int RESULT_WIDTH = 1280;

    private readonly Dictionary<Key, PlaybackControlEnum> _keysControlEnumDict = new()
    {
        { Key.Space, PlaybackControlEnum.Pause },
        { Key.Left, PlaybackControlEnum.PreviousFrame },
        { Key.Right, PlaybackControlEnum.NextFrame },
        { Key.Enter, PlaybackControlEnum.SaveAnnotations },
        { Key.Delete, PlaybackControlEnum.RemoveSelectedAnns },
        { Key.X, PlaybackControlEnum.RemoveAllAnns },
        { Key.PageUp, PlaybackControlEnum.Previous },
        { Key.PageDown, PlaybackControlEnum.Next },
    };

    public async Task Handle(AnnClassSelectedEvent notification, CancellationToken cancellationToken)
    {
        SelectClass(notification.DetectionClass);
        await Task.CompletedTask;
    }

    private void SelectClass(DetectionClass detClass)
    {
        mainWindow.Editor.CurrentAnnClass = detClass;
        foreach (var ann in mainWindow.Editor.CurrentDetections.Where(x => x.IsSelected))
            ann.DetectionClass = detClass;
        mainWindow.LvClasses.SelectNum(detClass.Id);
    }

    public async Task Handle(KeyEvent keyEvent, CancellationToken cancellationToken = default)
    {
        if (keyEvent.WindowEnum != WindowEnum.Annotator)
            return;

        var key = keyEvent.Args.Key;
        var keyNumber = (int?)null;

        if ((int)key >= (int)Key.D1 && (int)key <= (int)Key.D9)
            keyNumber = key - Key.D1;
        if ((int)key >= (int)Key.NumPad1 && (int)key <= (int)Key.NumPad9)
            keyNumber = key - Key.NumPad1;
        if (keyNumber.HasValue)
            SelectClass((DetectionClass)mainWindow.LvClasses.DetectionDataGrid.Items[keyNumber.Value]!);

        if (_keysControlEnumDict.TryGetValue(key, out var value))
            await ControlPlayback(value, cancellationToken);

        if (key == Key.R)
            mainWindow.AutoDetect(null!, null!);

        #region Volume
        switch (key)
        {
            case Key.VolumeMute when mediaPlayer.Volume == 0:
                await ControlPlayback(PlaybackControlEnum.TurnOnVolume, cancellationToken);
                break;
            case Key.VolumeMute:
                await ControlPlayback(PlaybackControlEnum.TurnOffVolume, cancellationToken);
                break;
            case Key.Up:
            case Key.VolumeUp:
                var vUp = Math.Min(100, mediaPlayer.Volume + 5);
                ChangeVolume(vUp);
                mainWindow.Volume.Value = vUp;
                break;
            case Key.Down:
            case Key.VolumeDown:
                var vDown = Math.Max(0, mediaPlayer.Volume - 5);
                ChangeVolume(vDown);
                mainWindow.Volume.Value = vDown;
                break;
        }
        #endregion
    }

    public async Task Handle(AnnotatorControlEvent notification, CancellationToken cancellationToken = default)
    {
        await ControlPlayback(notification.PlaybackControl, cancellationToken);
        mainWindow.VideoView.Focus();
    }

    private async Task ControlPlayback(PlaybackControlEnum controlEnum, CancellationToken cancellationToken = default)
    {
        try
        {
            var isCtrlPressed = Keyboard.IsKeyDown(Key.LeftCtrl) || Keyboard.IsKeyDown(Key.RightCtrl);
            var step = isCtrlPressed ? LARGE_STEP : STEP;

            switch (controlEnum)
            {
                case PlaybackControlEnum.Play:
                    await Play(cancellationToken);
                    break;
                case PlaybackControlEnum.Pause:
                    mediaPlayer.Pause();
                    if (mainWindow.IsInferenceNow)
                        mainWindow.FollowAI = false;
                    if (!mediaPlayer.IsPlaying)
                        mainWindow.BlinkHelp(HelpTexts.HelpTextsDict[HelpTextEnum.AnnotationHelp]);

                    if (formState.BackgroundTime.HasValue)
                    {
                        mainWindow.Editor.ResetBackground();
                        formState.BackgroundTime = null;
                    }
                    break;
                case PlaybackControlEnum.Stop:
                    inferenceService.StopInference();
                    await mainWindow.DetectionCancellationSource.CancelAsync();
                    mediaPlayer.Stop();
                    break;
                case PlaybackControlEnum.PreviousFrame:
                    mainWindow.SeekTo(mediaPlayer.Time - step);
                    break;
                case PlaybackControlEnum.NextFrame:
                    mainWindow.SeekTo(mediaPlayer.Time + step);
                    break;
                case PlaybackControlEnum.SaveAnnotations:
                    await SaveAnnotations(cancellationToken);
                    break;
                case PlaybackControlEnum.RemoveSelectedAnns:

                    mainWindow.Editor.RemoveSelectedAnns();
                    break;
                case PlaybackControlEnum.RemoveAllAnns:
                    mainWindow.Editor.RemoveAllAnns();
                    break;
                case PlaybackControlEnum.TurnOnVolume:
                    mainWindow.TurnOnVolumeBtn.Visibility = Visibility.Collapsed;
                    mainWindow.TurnOffVolumeBtn.Visibility = Visibility.Visible;
                    mediaPlayer.Volume = formState.CurrentVolume;
                    break;
                case PlaybackControlEnum.TurnOffVolume:
                    mainWindow.TurnOffVolumeBtn.Visibility = Visibility.Collapsed;
                    mainWindow.TurnOnVolumeBtn.Visibility = Visibility.Visible;
                    formState.CurrentVolume = mediaPlayer.Volume;
                    mediaPlayer.Volume = 0;
                    break;
                case PlaybackControlEnum.Previous:
                    await NextMedia(isPrevious: true, ct: cancellationToken);
                    break;
                case PlaybackControlEnum.Next:
                    await NextMedia(ct: cancellationToken);
                    break;
                case PlaybackControlEnum.None:
                    break;
                default:
                    throw new ArgumentOutOfRangeException(nameof(controlEnum), controlEnum, null);
            }
        }
        catch (Exception e)
        {
            logger.LogError(e, e.Message);
            throw;
        }
    }

    private async Task NextMedia(bool isPrevious = false, CancellationToken ct = default)
    {
        var increment = isPrevious ? -1 : 1;
        var check = isPrevious ? -1 : mainWindow.LvFiles.Items.Count;
        if (mainWindow.LvFiles.SelectedIndex + increment == check)
            return;

        mainWindow.LvFiles.SelectedIndex += increment;
        await Play(ct);
    }

    public async Task Handle(VolumeChangedEvent notification, CancellationToken cancellationToken)
    {
        ChangeVolume(notification.Volume);
        await Task.CompletedTask;
    }

    private void ChangeVolume(int volume)
    {
        formState.CurrentVolume = volume;
        mediaPlayer.Volume = volume;
    }

    private async Task Play(CancellationToken ct = default)
    {
        if (mainWindow.LvFiles.SelectedItem == null)
            return;
        var mediaInfo = (MediaFileInfo)mainWindow.LvFiles.SelectedItem;
        mainWindow.Editor.ResetBackground();

        formState.CurrentMedia = mediaInfo;
        //need to wait a bit for correct VLC playback event handling
        await Task.Delay(100, ct);
        mediaPlayer.Stop();
        mainWindow.Title = $"Azaion Annotator - {mediaInfo.Name}";
        mainWindow.BlinkHelp(HelpTexts.HelpTextsDict[HelpTextEnum.PauseForAnnotations]);
        mediaPlayer.Play(new Media(libVLC, mediaInfo.Path));
        if (formState.CurrentMedia.MediaType == MediaTypes.Image)
            mediaPlayer.SetPause(true);
    }

    //SAVE: MANUAL
    private async Task SaveAnnotations(CancellationToken cancellationToken = default)
    {
        if (formState.CurrentMedia == null)
            return;

        var time = formState.BackgroundTime ?? TimeSpan.FromMilliseconds(mediaPlayer.Time);
        var originalMediaName = formState.VideoName;
        var fName = originalMediaName.ToTimeName(time);

        var currentDetections = mainWindow.Editor.CurrentDetections
            .Select(x => new Detection(fName, x.GetLabel(mainWindow.Editor.RenderSize, formState.BackgroundTime.HasValue ? mainWindow.Editor.RenderSize : formState.CurrentVideoSize)))
            .ToList();

        formState.CurrentMedia.HasAnnotations = currentDetections.Count != 0;
        mainWindow.LvFiles.Items.Refresh();
        mainWindow.Editor.RemoveAllAnns();

        var isVideo = formState.CurrentMedia.MediaType == MediaTypes.Video;
        var imgPath = Path.Combine(dirConfig.Value.ImagesDirectory, $"{fName}{Constants.JPG_EXT}");

        if (formState.BackgroundTime.HasValue)
        {
            //no need to save image, it's already there, just remove background
            mainWindow.Editor.ResetBackground();
            formState.BackgroundTime = null;

            //next item
            var annGrid = mainWindow.DgAnnotations;
            annGrid.SelectedIndex = Math.Min(annGrid.Items.Count, annGrid.SelectedIndex + 1);
            mainWindow.OpenAnnotationResult((AnnotationResult)annGrid.SelectedItem);
        }
        else
        {
            var resultHeight = (uint)Math.Round(RESULT_WIDTH / formState.CurrentVideoSize.Width * formState.CurrentVideoSize.Height);
            mediaPlayer.TakeSnapshot(0, imgPath, RESULT_WIDTH, resultHeight);
            if (isVideo)
                mediaPlayer.Play();
            else
                await NextMedia(ct: cancellationToken);
        }

        var annotation = await annotationService.SaveAnnotation(originalMediaName, time, currentDetections, token: cancellationToken);
        if (isVideo)
            mainWindow.AddAnnotation(annotation);
    }

    public async Task Handle(AnnotationsDeletedEvent notification, CancellationToken cancellationToken)
    {
        var annResDict = formState.AnnotationResults.ToDictionary(x => x.Annotation.Name, x => x);
        foreach (var ann in notification.Annotations)
        {
            if (!annResDict.TryGetValue(ann.Name, out var value))
                continue;

            formState.AnnotationResults.Remove(value);
            mainWindow.TimedAnnotations.Remove(ann);
        }

        if (formState.AnnotationResults.Count == 0)
        {
            var media = mainWindow.AllMediaFiles.FirstOrDefault(x => x.Name == formState.CurrentMedia?.Name);
            if (media != null)
            {
                media.HasAnnotations = false;
                mainWindow.LvFiles.Items.Refresh();
            }
        }
        await Task.CompletedTask;
    }
}