annotations/Azaion.Dataset/DatasetExplorerEventHandler.cs

using System.Windows.Input;
using System.Windows.Media.Imaging;
using Azaion.Common.Database;
using Azaion.Common.DTO;
using Azaion.Common.Events;
using Azaion.Common.Services;
using MediatR;
using Microsoft.Extensions.Logging;

namespace Azaion.Dataset;

public class DatasetExplorerEventHandler(
    ILogger<DatasetExplorerEventHandler> logger,
    DatasetExplorer datasetExplorer,
    IAnnotationService annotationService,
    IAzaionApi azaionApi) :
        INotificationHandler<KeyEvent>,
        INotificationHandler<DatasetExplorerControlEvent>,
        INotificationHandler<AnnotationCreatedEvent>,
        INotificationHandler<AnnotationsDeletedEvent>
{
    private readonly Dictionary<Key, PlaybackControlEnum> _keysControlEnumDict = new()
    {
        { Key.Enter, PlaybackControlEnum.SaveAnnotations },
        { Key.Delete, PlaybackControlEnum.RemoveSelectedAnns },
        { Key.X, PlaybackControlEnum.RemoveAllAnns },
        { Key.Escape, PlaybackControlEnum.Close },
        { Key.Down, PlaybackControlEnum.Next },
        { Key.PageDown, PlaybackControlEnum.Next },
        { Key.Up, PlaybackControlEnum.Previous },
        { Key.PageUp, PlaybackControlEnum.Previous },
        { Key.V, PlaybackControlEnum.ValidateAnnotations},
    };

    public async Task Handle(DatasetExplorerControlEvent notification, CancellationToken token)
    {
        await HandleControl(notification.PlaybackControl, token);
    }

    public async Task Handle(KeyEvent keyEvent, CancellationToken token)
    {
        if (keyEvent.WindowEnum != WindowEnum.DatasetExplorer)
            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)
            datasetExplorer.LvClasses.SelectNum(keyNumber.Value);
        else
        {
            if (datasetExplorer.Switcher.SelectedIndex == 1 && _keysControlEnumDict.TryGetValue(key, out var value))
                await HandleControl(value, token);
        }
    }

    private async Task HandleControl(PlaybackControlEnum controlEnum, CancellationToken token = default)
    {
        switch (controlEnum)
        {
            case PlaybackControlEnum.SaveAnnotations:
                if (datasetExplorer.ThumbnailLoading)
                    return;

                var a = datasetExplorer.CurrentAnnotation!.Annotation;

                var mediaSize = ((BitmapImage)datasetExplorer.ExplorerEditor.BackgroundImage.Source).Size();
                var detections = datasetExplorer.ExplorerEditor.CurrentDetections
                    .Select(x => new Detection(a.Name, x.ToYoloLabel(datasetExplorer.ExplorerEditor.RenderSize, mediaSize)))
                    .ToList();
                var index = datasetExplorer.ThumbnailsView.SelectedIndex;
                var annotation = await annotationService.SaveAnnotation(a.OriginalMediaName, a.Name, a.Time, detections, token: token);
                await ValidateAnnotations([annotation], token);
                await datasetExplorer.EditAnnotation(index + 1);
                break;
            case PlaybackControlEnum.RemoveSelectedAnns:
                if (datasetExplorer.ExplorerEditor.CurrentDetections.Any(x => x.IsSelected))
                    datasetExplorer.ExplorerEditor.RemoveSelectedAnns();
                else
                {
                    await datasetExplorer.DeleteAnnotations();
                    await datasetExplorer.EditAnnotation(datasetExplorer.ThumbnailsView.SelectedIndex);
                }
                break;
            case PlaybackControlEnum.RemoveAllAnns:
                datasetExplorer.ExplorerEditor.RemoveAllAnns();
                break;
            case PlaybackControlEnum.Next:
                await datasetExplorer.EditAnnotation(datasetExplorer.ThumbnailsView.SelectedIndex + 1);
                break;
            case PlaybackControlEnum.Previous:
                await datasetExplorer.EditAnnotation(datasetExplorer.ThumbnailsView.SelectedIndex - 1);
                break;
            case PlaybackControlEnum.Close:
                datasetExplorer.SwitchTab(toEditor: false);
                datasetExplorer.SelectedAnnotationName.Text = "";
                break;
            case PlaybackControlEnum.ValidateAnnotations:
                var annotations = datasetExplorer.ThumbnailsView.SelectedItems.Cast<AnnotationThumbnail>()
                    .Select(x => x.Annotation)
                    .ToList();
                await ValidateAnnotations(annotations, token);
                break;
        }
    }

    private async Task ValidateAnnotations(List<Annotation> annotations, CancellationToken token = default)
    {
        await annotationService.ValidateAnnotations(annotations.Select(x => x.Name).ToList(), token: token);
        foreach (var ann in datasetExplorer.SelectedAnnotations.Where(x => annotations.Contains(x.Annotation)))
        {
            ann.Annotation.AnnotationStatus = AnnotationStatus.Validated;
            if (datasetExplorer.SelectedAnnotationDict.TryGetValue(ann.Annotation.Name, out var value))
                value.Annotation.AnnotationStatus = AnnotationStatus.Validated;
            ann.UpdateUI();
        }
    }

    public Task Handle(AnnotationCreatedEvent notification, CancellationToken token)
    {
        datasetExplorer.Dispatcher.Invoke(() =>
        {
            var annotation = notification.Annotation;
            var selectedClass = datasetExplorer.LvClasses.CurrentClassNumber;

            datasetExplorer.AddAnnotationToDict(annotation);
            if (!annotation.Classes.Contains(selectedClass) && selectedClass != -1)
                return;

            var index = 0;
            var annThumb = new AnnotationThumbnail(annotation, azaionApi.CurrentUser.Role.IsValidator());
            if (datasetExplorer.SelectedAnnotationDict.ContainsKey(annThumb.Annotation.Name))
            {
                datasetExplorer.SelectedAnnotationDict.Remove(annThumb.Annotation.Name);
                var ann = datasetExplorer.SelectedAnnotations.FirstOrDefault(x => x.Annotation.Name == annThumb.Annotation.Name);
                if (ann != null)
                {
                    index = datasetExplorer.SelectedAnnotations.IndexOf(ann);
                    datasetExplorer.SelectedAnnotations.Remove(ann);
                }
            }

            datasetExplorer.SelectedAnnotations.Insert(index, annThumb);
            datasetExplorer.SelectedAnnotationDict.Add(annThumb.Annotation.Name, annThumb);
        });
        return Task.CompletedTask;
    }

    public async Task Handle(AnnotationsDeletedEvent notification, CancellationToken token)
    {
        try
        {
            datasetExplorer.Dispatcher.Invoke(() =>
            {
                var annThumbs = datasetExplorer.SelectedAnnotationDict
                    .Where(x => notification.AnnotationNames.Contains(x.Key))
                    .Select(x => x.Value)
                    .ToList();
                foreach (var annThumb in annThumbs)
                {
                    datasetExplorer.SelectedAnnotations.Remove(annThumb);
                    datasetExplorer.SelectedAnnotationDict.Remove(annThumb.Annotation.Name);
                }
            });
        }
        catch (Exception e)
        {
            logger.LogError(e, e.Message);
            throw;
        }
        await Task.CompletedTask;
    }
}