annotations/Azaion.Dataset/DatasetExplorer.xaml.cs

using System.Collections.ObjectModel;
using System.Windows;
using System.Windows.Input;
using System.Windows.Media;
using Azaion.Common;
using Azaion.Common.Database;
using Azaion.Common.DTO;
using Azaion.Common.DTO.Config;
using Azaion.Common.Services;
using Azaion.CommonSecurity.DTO;
using Azaion.CommonSecurity.Services;
using LinqToDB;
using MediatR;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using ScottPlot;
using Color = ScottPlot.Color;

namespace Azaion.Dataset;

public partial class DatasetExplorer : INotificationHandler<AnnotationCreatedEvent>
{
    private readonly ILogger<DatasetExplorer> _logger;
    private readonly AnnotationConfig _annotationConfig;
    private readonly DirectoriesConfig _directoriesConfig;

    private Dictionary<int, List<Annotation>> _annotationsDict;

    public ObservableCollection<AnnotationImageView> SelectedAnnotations { get; set; } = new();
    public ObservableCollection<DetectionClass> AllAnnotationClasses { get; set; } = new();

    private Dictionary<string, LabelInfo> LabelsCache { get; set; } = new();

    private int _tempSelectedClassIdx = 0;
    private readonly IGalleryService _galleryService;
    private readonly IDbFactory _dbFactory;
    private readonly IMediator _mediator;
    private readonly AzaionApiClient _apiClient;

    public bool ThumbnailLoading { get; set; }

    public AnnotationImageView? CurrentAnnotation { get; set; }

    public DatasetExplorer(
        IOptions<DirectoriesConfig> directoriesConfig,
        IOptions<AnnotationConfig> annotationConfig,
        ILogger<DatasetExplorer> logger,
        IGalleryService galleryService,
        FormState formState,
        IDbFactory dbFactory,
        IMediator mediator,
        AzaionApiClient apiClient)
    {
        _directoriesConfig = directoriesConfig.Value;
        _annotationConfig = annotationConfig.Value;
        _logger = logger;
        _galleryService = galleryService;
        _dbFactory = dbFactory;
        _mediator = mediator;
        _apiClient = apiClient;

        InitializeComponent();
        Loaded += OnLoaded;
        Activated += (_, _) => formState.ActiveWindow = WindowEnum.DatasetExplorer;

         ThumbnailsView.KeyDown += async (sender, args) =>
        {
            switch (args.Key)
            {
                case Key.Delete:
                    DeleteAnnotations();
                    break;
                case Key.Enter:
                    await EditAnnotation();
                    break;
            }
        };
        ThumbnailsView.MouseDoubleClick += async (_, _) => await EditAnnotation();

        ThumbnailsView.SelectionChanged += (_, _) =>
        {
            StatusText.Text = $"Обрано: {ThumbnailsView.SelectedItems.Count} | {ThumbnailsView.SelectedIndex} / {SelectedAnnotations.Count}";
            ValidateBtn.Visibility = ThumbnailsView.SelectedItems.Cast<AnnotationImageView>().Any(x => x.IsSeed)
                ? Visibility.Visible
                : Visibility.Hidden;
        };
        ExplorerEditor.GetTimeFunc = () => Constants.GetTime(CurrentAnnotation!.Annotation.ImagePath);
    }

    private async void OnLoaded(object sender, RoutedEventArgs e)
    {
        AllAnnotationClasses = new ObservableCollection<DetectionClass>(
            new List<DetectionClass> { new() {Id = -1, Name = "All", ShortName = "All"}}
                .Concat(_annotationConfig.AnnotationClasses));
        LvClasses.ItemsSource = AllAnnotationClasses;

        LvClasses.MouseUp += async (_, _) =>
        {
            var selectedClass = (DetectionClass)LvClasses.SelectedItem;
            ExplorerEditor.CurrentAnnClass = selectedClass;
            _annotationConfig.LastSelectedExplorerClass = selectedClass.Id;

            if (Switcher.SelectedIndex == 0)
                await ReloadThumbnails();
            else
                foreach (var ann in ExplorerEditor.CurrentDetections.Where(x => x.IsSelected))
                    ann.DetectionClass = selectedClass;
        };

        LvClasses.SelectionChanged += (_, _) =>
        {
            if (Switcher.SelectedIndex != 1)
                return;

            var selectedClass = (DetectionClass)LvClasses.SelectedItem;
            if (selectedClass == null)
                return;

            ExplorerEditor.CurrentAnnClass = selectedClass;

            foreach (var ann in ExplorerEditor.CurrentDetections.Where(x => x.IsSelected))
                ann.DetectionClass = selectedClass;
        };

        LvClasses.SelectedIndex = _annotationConfig.LastSelectedExplorerClass ?? 0;
        ExplorerEditor.CurrentAnnClass = (DetectionClass)LvClasses.SelectedItem;

        await _dbFactory.Run(async db =>
        {
            var allAnnotations = await db.Annotations
                .LoadWith(x => x.Detections)
                .OrderByDescending(x => x.CreatedDate)
                .ToListAsync();
            _annotationsDict = AllAnnotationClasses.ToDictionary(x => x.Id, _ => new List<Annotation>());

            foreach (var annotation in allAnnotations)
                AddAnnotationToDict(annotation);
        });
        await ReloadThumbnails();
        await LoadClassDistribution();

        DataContext = this;
    }

    private void AddAnnotationToDict(Annotation annotation)
    {
        foreach (var c in annotation.Classes)
            _annotationsDict[c].Add(annotation);
        _annotationsDict[-1].Add(annotation);
    }

    private async Task LoadClassDistribution()
    {
        var data = LabelsCache
            .SelectMany(x => x.Value.Classes)
            .GroupBy(x => x)
            .Select(x => new
            {
                x.Key,
                _annotationConfig.DetectionClassesDict[x.Key].Name,
                _annotationConfig.DetectionClassesDict[x.Key].Color,
                ClassCount = x.Count()
            })
            .ToList();

        var foregroundColor = Color.FromColor(System.Drawing.Color.Black);

        var bars = data.Select(x => new Bar
        {
            Orientation = Orientation.Horizontal,
            Position = -1.5 * x.Key + 1,
            Label = x.ClassCount > 200 ? x.ClassCount.ToString() : "",
            FillColor = new Color(x.Color.R, x.Color.G, x.Color.B, x.Color.A),
            Value = x.ClassCount,
            CenterLabel = true,
            LabelOffset = 10
        }).ToList();

        ClassDistribution.Plot.Add.Bars(bars);

        foreach (var x in data)
        {
            var label = ClassDistribution.Plot.Add.Text(x.Name, 50, -1.5 * x.Key + 1.1);
            label.LabelFontColor = foregroundColor;
            label.LabelFontSize = 18;
        }

        ClassDistribution.Plot.Axes.AutoScale();
        ClassDistribution.Plot.HideAxesAndGrid();
        ClassDistribution.Plot.FigureBackground.Color = new("#888888");

        ClassDistribution.Refresh();
    }

    private void ReloadThumbnailsItemClick(object sender, RoutedEventArgs e)
    {
        var result = MessageBox.Show($"Видалити всі іконки і згенерувати нову базу іконок в {_directoriesConfig.ThumbnailsDirectory}?",
            "Підтвердження оновлення іконок", MessageBoxButton.YesNo, MessageBoxImage.Question);
        if (result != MessageBoxResult.Yes)
            return;
        _galleryService.ClearThumbnails();
        _galleryService.RefreshThumbnails();
    }

    private async Task EditAnnotation()
    {
        try
        {
            ThumbnailLoading = true;

            if (ThumbnailsView.SelectedItem == null)
                return;

            CurrentAnnotation = (ThumbnailsView.SelectedItem as AnnotationImageView)!;
            var ann = CurrentAnnotation.Annotation;
            ExplorerEditor.Background = new ImageBrush
            {
                ImageSource = await ann.ImagePath.OpenImage()
            };
            SwitchTab(toEditor: true);

            var time = Constants.GetTime(ann.ImagePath);
            ExplorerEditor.RemoveAllAnns();
            foreach (var deetection in ann.Detections)
            {
                var annClass = _annotationConfig.DetectionClassesDict[deetection.ClassNumber];
                var canvasLabel = new CanvasLabel(deetection, ExplorerEditor.RenderSize, ExplorerEditor.RenderSize);
                ExplorerEditor.CreateAnnotation(annClass, time, canvasLabel);
            }

            ThumbnailLoading = false;
        }
        catch (Exception e)
        {
            _logger.LogError(e, e.Message);
            throw;
        }
        finally
        {
            ThumbnailLoading = false;
        }

    }

    public void SwitchTab(bool toEditor)
    {
        if (toEditor)
        {
            AnnotationsTab.Visibility = Visibility.Collapsed;
            EditorTab.Visibility = Visibility.Visible;
            _tempSelectedClassIdx = LvClasses.SelectedIndex;
            LvClasses.ItemsSource = _annotationConfig.AnnotationClasses;

            Switcher.SelectedIndex = 1;
            LvClasses.SelectedIndex = Math.Max(0, _tempSelectedClassIdx - 1);
        }
        else
        {
            AnnotationsTab.Visibility = Visibility.Visible;
            EditorTab.Visibility = Visibility.Collapsed;
            LvClasses.ItemsSource = AllAnnotationClasses;
            LvClasses.SelectedIndex = _tempSelectedClassIdx;
            Switcher.SelectedIndex = 0;
        }
    }

    private void DeleteAnnotations()
    {
        var tempSelected = ThumbnailsView.SelectedIndex;
        var result = MessageBox.Show("Чи дійсно видалити аннотації?","Підтвердження видалення", MessageBoxButton.YesNo, MessageBoxImage.Question);
        if (result != MessageBoxResult.Yes)
            return;

        var selected = ThumbnailsView.SelectedItems.Count;
        for (var i = 0; i < selected; i++)
        {
            var dto = (ThumbnailsView.SelectedItems[0] as AnnotationImageView)!;
            dto.Delete();
            SelectedAnnotations.Remove(dto);
        }
        ThumbnailsView.SelectedIndex = Math.Min(SelectedAnnotations.Count, tempSelected);
    }

    private async Task ReloadThumbnails()
    {
        SelectedAnnotations.Clear();
        foreach (var ann in _annotationsDict[ExplorerEditor.CurrentAnnClass.Id])
            SelectedAnnotations.Add(new AnnotationImageView(ann));
    }

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

        //TODO: For editing existing need to handle updates
        AddAnnotationToDict(annotation);
        if (annotation.Classes.Contains(selectedClass.Value))
        {
            SelectedAnnotations.Add(new AnnotationImageView(annotation));
        }
    }

    private async void ValidateAnnotationsClick(object sender, RoutedEventArgs e)
    {
        try
        {
            await _mediator.Publish(new DatasetExplorerControlEvent(PlaybackControlEnum.ValidateAnnotations));
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, ex.Message);
        }
    }
}