annotations/Azaion.Annotator/AnnotatorEventHandler.cs

using System.IO;
using System.Windows;
using System.Windows.Input;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using Azaion.Annotator.Controls;
using Azaion.Annotator.DTO;
using Azaion.Common;
using Azaion.Common.Database;
using Azaion.Common.DTO;
using Azaion.Common.DTO.Config;
using Azaion.Common.Events;
using Azaion.Common.Extensions;
using Azaion.Common.Services;
using Azaion.Common.Services.Inference;
using GMap.NET;
using GMap.NET.WindowsPresentation;
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,
    IAnnotationService annotationService,
    ILogger<AnnotatorEventHandler> logger,
    IOptions<DirectoriesConfig> dirConfig,
    IOptions<AnnotationConfig> annotationConfig,
    IInferenceService inferenceService,
    IDbFactory dbFactory,
    IAzaionApi api,
    FailsafeAnnotationsProducer producer)
    :
            INotificationHandler<KeyEvent>,
            INotificationHandler<AnnClassSelectedEvent>,
            INotificationHandler<AnnotatorControlEvent>,
            INotificationHandler<VolumeChangedEvent>,
            INotificationHandler<AnnotationsDeletedEvent>,
            INotificationHandler<AnnotationAddedEvent>,
            INotificationHandler<SetStatusTextEvent>,
            INotificationHandler<GPSMatcherResultProcessedEvent>,
            INotificationHandler<AIAvailabilityStatusEvent>
{
    private const int STEP = 20;
    private const int LARGE_STEP = 5000;
    private readonly string _tempImgPath = Path.Combine(dirConfig.Value.ImagesDirectory, "___temp___.jpg"); 
    
    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 ct)
    {
        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 ct = 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, ct);

        if (key == Key.R)
            await mainWindow.AutoDetect();

        #region Volume
        switch (key)
        {
            case Key.VolumeMute when mediaPlayer.Volume == 0:
                await ControlPlayback(PlaybackControlEnum.TurnOnVolume, ct);
                break;
            case Key.VolumeMute:
                await ControlPlayback(PlaybackControlEnum.TurnOffVolume, ct);
                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 ct = default)
    {
        await ControlPlayback(notification.PlaybackControl, ct);
        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:
                    if (mediaPlayer.IsPlaying) 
                    {
                        mediaPlayer.Pause();
                        mediaPlayer.TakeSnapshot(0, _tempImgPath, 0, 0);
                        mainWindow.Editor.SetBackground(await _tempImgPath.OpenImage());
                        formState.BackgroundTime = TimeSpan.FromMilliseconds(mediaPlayer.Time);
                    }
                    else
                    {
                        mediaPlayer.Play();
                        if (formState.BackgroundTime.HasValue)
                        {
                            mainWindow.Editor.SetBackground(null);
                            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 SaveAnnotation(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 ct)
    {
        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;
        
        if (formState.CurrentMedia == mediaInfo)
            return; //already loaded
        
        formState.CurrentMedia = mediaInfo;
        mainWindow.Title =  $"{mainWindow.MainTitle} - {mediaInfo.Name}";
            
        await mainWindow.ReloadAnnotations();
            
        if (mediaInfo.MediaType == MediaTypes.Video)
        {
            mainWindow.Editor.SetBackground(null);
            //need to wait a bit for correct VLC playback event handling
            await Task.Delay(100, ct);
            mediaPlayer.Stop();
            mediaPlayer.Play(new Media(libVlc, mediaInfo.Path));   
        }
        else
        {
            formState.BackgroundTime = TimeSpan.Zero;
            var image = await mediaInfo.Path.OpenImage();
            formState.CurrentMediaSize = new Size(image.PixelWidth, image.PixelHeight);
            mainWindow.Editor.SetBackground(image);
            mediaPlayer.Stop();
            mainWindow.ShowTimeAnnotations(TimeSpan.Zero, showImage: true);
        }
    }

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

        var time = formState.BackgroundTime ?? TimeSpan.FromMilliseconds(mediaPlayer.Time);
        var timeName = formState.MediaName.ToTimeName(time);
        var isVideo = formState.CurrentMedia.MediaType == MediaTypes.Video;
        var imgPath = Path.Combine(dirConfig.Value.ImagesDirectory, $"{timeName}{Constants.JPG_EXT}");

        formState.CurrentMedia.HasAnnotations = mainWindow.Editor.CurrentDetections.Count != 0;
        var annotations = await SaveAnnotationInner(imgPath, cancellationToken);
        if (isVideo)
        {
            foreach (var annotation in annotations) 
                mainWindow.AddAnnotation(annotation);
            mediaPlayer.Play();
            
            // next item. Probably not needed
            // var annGrid = mainWindow.DgAnnotations;
            // annGrid.SelectedIndex = Math.Min(annGrid.Items.Count, annGrid.SelectedIndex + 1);
            // mainWindow.OpenAnnotationResult((AnnotationResult)annGrid.SelectedItem);
            
            mainWindow.Editor.SetBackground(null);
            formState.BackgroundTime = null;
        }
        else
        {
            await NextMedia(ct: cancellationToken);
        }
        
        mainWindow.LvFiles.Items.Refresh();
        mainWindow.Editor.RemoveAllAnns();
    }

    private async Task<List<Annotation>> SaveAnnotationInner(string imgPath, CancellationToken cancellationToken = default)
    {
        var canvasDetections = mainWindow.Editor.CurrentDetections.Select(x => x.ToCanvasLabel()).ToList();
        var annotationsResult = new List<Annotation>();
        if (!File.Exists(imgPath))
        {
            var source = (mainWindow.Editor.BackgroundImage.Source as BitmapSource)!;
            if (new Size(source.PixelWidth, source.PixelHeight).FitSizeForAI())
                await source.SaveImage(imgPath, cancellationToken);
            else
            {
                //Tiling
                
                //1. Convert from RenderSize to CurrentMediaSize
                var detectionCoords = canvasDetections.Select(x => new CanvasLabel(
                    new YoloLabel(x, mainWindow.Editor.RenderSize, formState.CurrentMediaSize), formState.CurrentMediaSize, null, x.Confidence))
                    .ToList();
                
                //2. Split to frames
                var results = TileProcessor.Split(formState.CurrentMediaSize, detectionCoords, cancellationToken);
                
                //3. Save each frame as a separate annotation
                foreach (var res in results)
                {
                    var time = TimeSpan.Zero;
                    var annotationName = $"{formState.MediaName}{Constants.SPLIT_SUFFIX}{res.Tile.Width}{res.Tile.Left:0000}_{res.Tile.Top:0000}!".ToTimeName(time);
                    
                    var tileImgPath = Path.Combine(dirConfig.Value.ImagesDirectory, $"{annotationName}{Constants.JPG_EXT}");
                    var bitmap = new CroppedBitmap(source, new Int32Rect((int)res.Tile.Left, (int)res.Tile.Top, (int)res.Tile.Width, (int)res.Tile.Height));
                    await bitmap.SaveImage(tileImgPath, cancellationToken);

                    var frameSize = new Size(res.Tile.Width, res.Tile.Height);
                    var detections = res.Detections
                        .Select(det => det.ReframeToSmall(res.Tile))
                        .Select(x => new Detection(annotationName, new YoloLabel(x, frameSize)))
                        .ToList();

                    annotationsResult.Add(await annotationService.SaveAnnotation(formState.MediaName, annotationName, time, detections, token: cancellationToken));
                }
                return annotationsResult;
            }
        }
        
        var timeImg = formState.BackgroundTime ?? TimeSpan.FromMilliseconds(mediaPlayer.Time);
        var annName = formState.MediaName.ToTimeName(timeImg);
        var currentDetections = canvasDetections.Select(x => 
                new Detection(annName, new YoloLabel(x, mainWindow.Editor.RenderSize)))
            .ToList();
        var annotation = await annotationService.SaveAnnotation(formState.MediaName, annName, timeImg, currentDetections, token: cancellationToken);
        return [annotation];
    }

    public async Task Handle(AnnotationsDeletedEvent notification, CancellationToken ct)
    {
        try
        {
            mainWindow.Dispatcher.Invoke(() =>
            {
                var namesSet = notification.AnnotationNames.ToHashSet();

                var remainAnnotations = formState.AnnotationResults
                    .Where(x => !namesSet.Contains(x.Name)).ToList();
                formState.AnnotationResults.Clear();
                foreach (var ann in remainAnnotations)
                    formState.AnnotationResults.Add(ann);

                var timedAnnotationsToRemove = mainWindow.TimedAnnotations
                    .Where(x => namesSet.Contains(x.Value.Name))
                    .Select(x => x.Value).ToList();
                mainWindow.TimedAnnotations.Remove(timedAnnotationsToRemove);

                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 dbFactory.DeleteAnnotations(notification.AnnotationNames, ct);
            
            foreach (var name in notification.AnnotationNames)
            {
                try
                {
                    File.Delete(Path.Combine(dirConfig.Value.ImagesDirectory, $"{name}{Constants.JPG_EXT}"));
                    File.Delete(Path.Combine(dirConfig.Value.LabelsDirectory, $"{name}{Constants.TXT_EXT}"));
                    File.Delete(Path.Combine(dirConfig.Value.ThumbnailsDirectory, $"{name}{Constants.THUMBNAIL_PREFIX}{Constants.JPG_EXT}"));
                    File.Delete(Path.Combine(dirConfig.Value.ResultsDirectory, $"{name}{Constants.RESULT_PREFIX}{Constants.JPG_EXT}"));
                }
                catch (Exception e)
                {
                    logger.LogError(e, e.Message);
                }
            }

            //Only validators can send Delete to the queue
            if (!notification.FromQueue && api.CurrentUser.Role.IsValidator())
                await producer.SendToInnerQueue(notification.AnnotationNames, AnnotationStatus.Deleted, ct);
        }
        catch (Exception e)
        {
            logger.LogError(e, e.Message);
            throw;
        }
    }

    public Task Handle(AnnotationAddedEvent e, CancellationToken ct)
    {
        mainWindow.Dispatcher.Invoke(() =>
        {
            
            var mediaInfo = (MediaFileInfo)mainWindow.LvFiles.SelectedItem;
            if ((mediaInfo?.FName ?? "") == e.Annotation.OriginalMediaName)
                mainWindow.AddAnnotation(e.Annotation);

            var log = string.Join(Environment.NewLine, e.Annotation.Detections.Select(det =>
                $"Розпізнавання {e.Annotation.OriginalMediaName}: {annotationConfig.Value.DetectionClassesDict[det.ClassNumber].ShortName}: " +
                $"xy=({det.CenterX:F2},{det.CenterY:F2}), " +
                $"розмір=({det.Width:F2}, {det.Height:F2}), " +
                $"conf: {det.Confidence*100:F0}%"));

            mainWindow.LvFiles.Items.Refresh();

            var media = mainWindow.MediaFilesDict.GetValueOrDefault(e.Annotation.OriginalMediaName);
            if (media != null)
                media.HasAnnotations = true;

            mainWindow.LvFiles.Items.Refresh();
            mainWindow.StatusHelp.Text = log;
        });
        return Task.CompletedTask;
    }

    public Task Handle(SetStatusTextEvent e, CancellationToken cancellationToken)
    {
        mainWindow.Dispatcher.Invoke(() =>
        {
            mainWindow.StatusHelp.Text = e.Text;
            mainWindow.StatusHelp.Foreground = e.IsError ? Brushes.Red : Brushes.White;
        });
        return Task.CompletedTask;
    }

    public Task Handle(GPSMatcherResultProcessedEvent e, CancellationToken cancellationToken)
    {
        mainWindow.Dispatcher.Invoke(() =>
        {
            var ann = mainWindow.MapMatcherComponent.Annotations[e.Index];
            AddMarker(e.GeoPoint, e.Image, Brushes.Blue);
            if (e.ProcessedGeoPoint != e.GeoPoint)
                AddMarker(e.ProcessedGeoPoint, $"{e.Image}: corrected", Brushes.DarkViolet);
            ann.Lat = e.GeoPoint.Lat;
            ann.Lon = e.GeoPoint.Lon;

        });
        return Task.CompletedTask;
    }

    private void AddMarker(GeoPoint point, string text, SolidColorBrush color)
    {
        var map = mainWindow.MapMatcherComponent;
        var pointLatLon = new PointLatLng(point.Lat, point.Lon);
        var marker = new GMapMarker(pointLatLon);
        marker.Shape = new CircleVisual(marker, size: 14, text: text, background: color);
        map.SatelliteMap.Markers.Add(marker);
        map.SatelliteMap.Position = pointLatLon;
        map.SatelliteMap.ZoomAndCenterMarkers(null);
    }

    public async Task Handle(AIAvailabilityStatusEvent e, CancellationToken cancellationToken)
    {
        mainWindow.Dispatcher.Invoke(() =>
        {
            logger.LogInformation(e.ToString());
            mainWindow.AIDetectBtn.IsEnabled = e.Status == AIAvailabilityEnum.Enabled;
            mainWindow.StatusHelp.Text = e.ToString();
        });
        if (e.Status is AIAvailabilityEnum.Enabled or AIAvailabilityEnum.Error)
            await inferenceService.CheckAIAvailabilityTokenSource.CancelAsync();
    }
}