annotations/Azaion.Dataset/DatasetExplorerEventHandler.cs

using System.IO;
using System.Windows.Input;
using Azaion.Common.DTO;
using Azaion.Common.DTO.Queue;
using Azaion.Common.Events;
using Azaion.Common.Services;
using MediatR;

namespace Azaion.Dataset;

public class DatasetExplorerEventHandler(
    DatasetExplorer datasetExplorer,
    AnnotationService annotationService) :
        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.V, PlaybackControlEnum.ValidateAnnotations}
    };

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

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

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

                var fName = Path.GetFileNameWithoutExtension(datasetExplorer.CurrentAnnotation!.Annotation.ImagePath);
                var extension = Path.GetExtension(fName);

                var detections = datasetExplorer.ExplorerEditor.CurrentDetections
                    .Select(x => new Detection(fName, x.GetLabel(datasetExplorer.ExplorerEditor.RenderSize)))
                    .ToList();
                await annotationService.SaveAnnotation(fName, extension, detections, SourceEnum.Manual, token: cancellationToken);
                datasetExplorer.SwitchTab(toEditor: false);
                break;
            case PlaybackControlEnum.RemoveSelectedAnns:
                datasetExplorer.ExplorerEditor.RemoveSelectedAnns();
                break;
            case PlaybackControlEnum.RemoveAllAnns:
                datasetExplorer.ExplorerEditor.RemoveAllAnns();
                break;
            case PlaybackControlEnum.Close:
                datasetExplorer.SwitchTab(toEditor: false);
                break;
            case PlaybackControlEnum.ValidateAnnotations:
                var annotations = datasetExplorer.ThumbnailsView.SelectedItems.Cast<AnnotationThumbnail>()
                    .Select(x => x.Annotation)
                    .ToList();
                foreach (var annotation in annotations)
                    await annotationService.ValidateAnnotation(annotation, cancellationToken);
                break;
        }
    }

    public async Task Handle(AnnotationCreatedEvent notification, CancellationToken cancellationToken)
    {
        var annotation = notification.Annotation;
        var selectedClass = ((DetectionClass?)datasetExplorer.LvClasses.SelectedItem)?.Id;
        if (selectedClass == null)
            return;

        //TODO: For editing existing need to handle updates
        datasetExplorer.AddAnnotationToDict(annotation);
        if (annotation.Classes.Contains(selectedClass.Value))
        {
            var annThumb = new AnnotationThumbnail(annotation);
            datasetExplorer.SelectedAnnotations.Add(annThumb);
            datasetExplorer.SelectedAnnotationDict.Add(annThumb.Annotation.Name, annThumb);
        }
    }

    public async Task Handle(AnnotationsDeletedEvent notification, CancellationToken cancellationToken)
    {
        var names = notification.Annotations.Select(x => x.Name).ToList();
        var annThumbs = datasetExplorer.SelectedAnnotationDict
            .Where(x => names.Contains(x.Key))
            .Select(x => x.Value)
            .ToList();
        foreach (var annThumb in annThumbs)
        {
            datasetExplorer.SelectedAnnotations.Remove(annThumb);
            datasetExplorer.SelectedAnnotationDict.Remove(annThumb.Annotation.Name);
        }
    }
}