annotations/Azaion.Dataset/DatasetExplorer.xaml.cs

using System.Collections.ObjectModel;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Input;
using System.Windows.Media;
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 LinqToDB;
using MediatR;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;

namespace Azaion.Dataset;

public partial class DatasetExplorer
{
    private readonly ILogger<DatasetExplorer> _logger;
    private readonly AppConfig _appConfig;

    private readonly Dictionary<int, Dictionary<string, Annotation>> _annotationsDict;
    private readonly CancellationTokenSource _cts = new();

    private List<DetectionClass> AllDetectionClasses { get; }
    public ObservableCollection<AnnotationThumbnail> SelectedAnnotations { get; } = new();
    public readonly Dictionary<string, AnnotationThumbnail> SelectedAnnotationDict = new();

    private int _tempSelectedClassIdx;
    private readonly IGalleryService _galleryService;
    private readonly IDbFactory _dbFactory;
    private readonly IMediator _mediator;

    private readonly IAzaionApi _azaionApi;
    private readonly IConfigUpdater _configUpdater;

    public bool ThumbnailLoading { get; set; }

    public AnnotationThumbnail? CurrentAnnotation { get; set; }

    private static readonly Guid SearchActionId = Guid.NewGuid();
    
    public DatasetExplorer(
        IOptions<AppConfig> appConfig,
        ILogger<DatasetExplorer> logger,
        IGalleryService galleryService,
        FormState formState,
        IDbFactory dbFactory,
        IMediator mediator,
        IAzaionApi azaionApi,
        IConfigUpdater configUpdater)
    {
        InitializeComponent();
        _appConfig = appConfig.Value;
        _logger = logger;
        _galleryService = galleryService;
        _dbFactory = dbFactory;
        _mediator = mediator;
        _azaionApi = azaionApi;
        _configUpdater = configUpdater;

        ShowWithObjectsOnlyChBox.IsChecked = _appConfig.UIConfig.ShowDatasetWithDetectionsOnly;
        var photoModes = Enum.GetValues(typeof(PhotoMode)).Cast<PhotoMode>().ToList();
        _annotationsDict = _appConfig.AnnotationConfig.DetectionClasses.SelectMany(cls => photoModes.Select(mode => (int)mode + cls.Id))
            .ToDictionary(x => x, _ => new Dictionary<string, Annotation>());
        _annotationsDict.Add(-1, []);

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

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

        ThumbnailsView.SelectionChanged += (_, _) =>
        {
            StatusText.Text = $"Обрано: {ThumbnailsView.SelectedItems.Count} | {ThumbnailsView.SelectedIndex} / {SelectedAnnotations.Count}";

            ValidateBtn.Visibility = ThumbnailsView.SelectedItems.Cast<AnnotationThumbnail>().Any(x => x.IsSeed)
                ? Visibility.Visible
                : Visibility.Hidden;
        };
        ExplorerEditor.GetTimeFunc = () => CurrentAnnotation!.Annotation.Time;
        _galleryService.ThumbnailsUpdate += thumbnailsPercentage =>
        {
            Dispatcher.Invoke(() => RefreshThumbBar.Value = thumbnailsPercentage);
        };
        Closing += (_, _) => _cts.Cancel();

        AllDetectionClasses = new List<DetectionClass>(
            new List<DetectionClass> { new() {Id = -1, Name = "All", ShortName = "All"}}
                .Concat(_appConfig.AnnotationConfig.DetectionClasses));
        LvClasses.Init(AllDetectionClasses);
    }

    private async void OnLoaded(object sender, RoutedEventArgs e)
    {
        LvClasses.DetectionClassChanged += async (_, args) =>
        {
            ExplorerEditor.CurrentAnnClass = args.DetectionClass;

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

        ExplorerEditor.CurrentAnnClass = LvClasses.CurrentDetectionClass ?? _appConfig.AnnotationConfig.DetectionClasses.First();

        var allAnnotations = await _dbFactory.Run(async db =>
            await db.Annotations.LoadWith(x => x.Detections)
                .OrderBy(x => x.AnnotationStatus)
                .ThenByDescending(x => x.CreatedDate)
                .ToListAsync());

        foreach (var annotation in allAnnotations)
            AddAnnotationToDict(annotation);

        await ReloadThumbnails();
        LoadClassDistribution();

        DataContext = this;
    }

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

    private void LoadClassDistribution()
    {
        var data = _annotationsDict
            .Where(x => x.Key != -1)
            .OrderBy(x => x.Key)
            .Select(gr => new ClusterDistribution
            {
                Label = $"{_appConfig.AnnotationConfig.DetectionClassesDict[gr.Key].UIName}: {gr.Value.Count}",
                Color = _appConfig.AnnotationConfig.DetectionClassesDict[gr.Key].Color,
                ClassCount = gr.Value.Count
            })
            .Where(x => x.ClassCount > 0)
            .ToList();

        var maxClassCount = Math.Max(1, data.Max(x => x.ClassCount));

        foreach (var cl in data)
        {
            cl.Color = cl.Color.CreateTransparent(150);
            cl.BarWidth = Math.Clamp(cl.ClassCount / (double)maxClassCount, 0, 1);
        }

        ClassDistributionPlot.Items = data;
    }

    private async void RefreshThumbnailsBtnClick(object sender, RoutedEventArgs e)
    {
        RefreshThumbnailsButtonItem.Visibility = Visibility.Hidden;
        RefreshProgressBarItem.Visibility = Visibility.Visible;

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

        RefreshProgressBarItem.Visibility = Visibility.Hidden;
        RefreshThumbnailsButtonItem.Visibility = Visibility.Visible;
    }

    public async Task EditAnnotation(int index)
    {
        try
        {
            ThumbnailLoading = true;
            if (index == -1)
                return;

            CurrentAnnotation = (ThumbnailsView.Items[index] as AnnotationThumbnail)!;
            ThumbnailsView.SelectedIndex = index;

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

            ExplorerEditor.RemoveAllAnns();
            ExplorerEditor.CreateDetections(ann, _appConfig.AnnotationConfig.DetectionClasses, ExplorerEditor.RenderSize);
        }
        catch (Exception e)
        {
            _logger.LogError(e, e.Message);
            throw;
        }
        finally
        {
            _ = Task.Run(async () =>
            {
                await Task.Delay(100);
                ThumbnailLoading = false;
            });
        }

    }

    public void SwitchTab(bool toEditor)
    {
        if (toEditor)
        {
            AnnotationsTab.Visibility = Visibility.Collapsed;
            EditorTab.Visibility = Visibility.Visible;
            _tempSelectedClassIdx = LvClasses.CurrentClassNumber;
            LvClasses.DetectionDataGrid.ItemsSource = _appConfig.AnnotationConfig.DetectionClasses;

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

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

        var annotationNames = ThumbnailsView.SelectedItems.Cast<AnnotationThumbnail>().Select(x => x.Annotation.Name).ToList();

        await _mediator.Publish(new AnnotationsDeletedEvent(annotationNames));
        ThumbnailsView.SelectedIndex = Math.Min(SelectedAnnotations.Count, tempSelected);
    }

    private async Task ReloadThumbnails()
    {
        var withDetectionsOnly = ShowWithObjectsOnlyChBox.IsChecked;
        SelectedAnnotations.Clear();
        SelectedAnnotationDict.Clear();
        var annThumbnails = _annotationsDict[ExplorerEditor.CurrentAnnClass.YoloId]
            .WhereIf(withDetectionsOnly, x => x.Value.Detections.Any())
            .WhereIf(TbSearch.Text.Length > 2, x => x.Key.ToLower().Contains(TbSearch.Text))
            .Select(x => new AnnotationThumbnail(x.Value, _azaionApi.CurrentUser.Role.IsValidator()))
            .OrderBy(x => !x.IsSeed)
            .ThenByDescending(x =>x.Annotation.CreatedDate);

        foreach (var thumb in annThumbnails)
        {
            SelectedAnnotations.Add(thumb);
            SelectedAnnotationDict.Add(thumb.Annotation.Name, thumb);
        }
        await Task.CompletedTask;
    }

    private async void ValidateAnnotationsClick(object sender, RoutedEventArgs e)
    {
        var result = MessageBox.Show("Підтверджуєте валідність обраних аннотацій?","Підтвердження валідності", MessageBoxButton.OKCancel, MessageBoxImage.Question);
        if (result != MessageBoxResult.OK)
            return;

        try
        {
            await _mediator.Publish(new DatasetExplorerControlEvent(PlaybackControlEnum.ValidateAnnotations), _cts.Token);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, ex.Message);
        }
    }

    private async void ShowWithObjectsOnly_OnClick(object sender, RoutedEventArgs e)
    {
        _appConfig.UIConfig.ShowDatasetWithDetectionsOnly = (sender as CheckBox)?.IsChecked ?? false;
        _configUpdater.Save(_appConfig);
        await ReloadThumbnails();
    }
    
    private void TbSearch_OnTextChanged(object sender, TextChangedEventArgs e)
    {
        TbSearch.Foreground = TbSearch.Text.Length > 2 ? Brushes.Black : Brushes.Gray;
        ThrottleExt.Throttle(ReloadThumbnails, SearchActionId, TimeSpan.FromMilliseconds(400));;
    }
}