annotations/Azaion.Annotator/DatasetExplorer.xaml.cs

using System.Collections.ObjectModel;
using System.IO;
using System.Windows;
using System.Windows.Input;
using System.Windows.Media;
using Azaion.Annotator.DTO;
using Azaion.Annotator.Extensions;
using Microsoft.Extensions.Logging;
using ScottPlot;
using Color = ScottPlot.Color;
using MessageBox = System.Windows.MessageBox;
using Orientation = ScottPlot.Orientation;

namespace Azaion.Annotator;

public partial class DatasetExplorer
{
    private readonly Config _config;
    private readonly ILogger<DatasetExplorer> _logger;

    public ObservableCollection<ThumbnailDto> ThumbnailsDtos { get; set; } = new();
    private ObservableCollection<AnnotationClass> AllAnnotationClasses { get; set; } = new();

    private int _tempSelectedClassIdx = 0;
    private readonly IConfigRepository _configRepository;
    private readonly FormState _formState;
    private readonly IGalleryManager _galleryManager;

    public bool ThumbnailLoading { get; set; }

    public ThumbnailDto? CurrentThumbnail { get; set; }

    public DatasetExplorer(
        Config config,
        ILogger<DatasetExplorer> logger,
        IConfigRepository configRepository,
        FormState formState,
        IGalleryManager galleryManager)
    {
        _config = config;
        _logger = logger;
        _configRepository = configRepository;
        _formState = formState;
        _galleryManager = galleryManager;

        InitializeComponent();
        Loaded += async (_, _) =>
        {
            AllAnnotationClasses = new ObservableCollection<AnnotationClass>(
                new List<AnnotationClass> { new() {Id = -1, Name = "All", ShortName = "All"}}
                    .Concat(_config.AnnotationClasses));
            LvClasses.ItemsSource = AllAnnotationClasses;

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

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

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

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

                ExplorerEditor.CurrentAnnClass = selectedClass;

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


            LvClasses.SelectedIndex = config.LastSelectedExplorerClass ?? 0;
            ExplorerEditor.CurrentAnnClass = (AnnotationClass)LvClasses.SelectedItem;
            await ReloadThumbnails();
            LoadClassDistribution();

            SizeChanged += async (_, _) => await SaveUserSettings();
            LocationChanged += async (_, _) => await SaveUserSettings();
            StateChanged += async (_, _) => await SaveUserSettings();

            RefreshThumbBar.Value = galleryManager.ThumbnailsPercentage;
            DataContext = this;
        };

        Closing += (sender, args) =>
        {
            args.Cancel = true;
            Visibility = Visibility.Hidden;
        };

        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} / {ThumbnailsDtos.Count}";
        };

        Activated += (_, _) => { _formState.ActiveWindow = WindowsEnum.DatasetExplorer; };

        ExplorerEditor.GetTimeFunc = () => _formState.GetTime(CurrentThumbnail!.ImagePath);
        galleryManager.ThumbnailsUpdate += thumbnailsPercentage =>
        {
            Dispatcher.Invoke(() => RefreshThumbBar.Value = thumbnailsPercentage);
        };

    }

    private void LoadClassDistribution()
    {
        var data = _galleryManager.LabelsCache
            .SelectMany(x => x.Value.Classes)
            .GroupBy(x => x)
            .Select(x => new
            {
                x.Key,
                _config.AnnotationClassesDict[x.Key].Name,
                _config.AnnotationClassesDict[x.Key].Color,
                ClassCount = x.Count()
            })
            .ToList();

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

        plot.Add.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
        }));

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

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

        ClassDistribution.Refresh();
    }

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

            if (ThumbnailsView.SelectedItem == null)
                return;

            var dto = (ThumbnailsView.SelectedItem as ThumbnailDto)!;
            CurrentThumbnail = dto;
            ExplorerEditor.Background = new ImageBrush
            {
                ImageSource = await dto.ImagePath.OpenImage()
            };
            SwitchTab(toEditor: true);

            var time = _formState.GetTime(dto.ImagePath);
            ExplorerEditor.RemoveAllAnns();
            foreach (var ann in await YoloLabel.ReadFromFile(dto.LabelPath))
            {
                var annClass = _config.AnnotationClassesDict[ann.ClassNumber];
                var canvasLabel = new CanvasLabel(ann, 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 = _config.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 async Task SaveUserSettings()
    {
        _config.DatasetExplorerConfig = this.GetConfig();
        await ThrottleExt.Throttle(() =>
        {
            _configRepository.Save(_config);
            return Task.CompletedTask;
        }, TimeSpan.FromSeconds(5));
    }

    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 ThumbnailDto)!;
            File.Delete(dto.ImagePath);
            File.Delete(dto.LabelPath);
            File.Delete(dto.ThumbnailPath);
            ThumbnailsDtos.Remove(dto);
        }
        ThumbnailsView.SelectedIndex = Math.Min(ThumbnailsDtos.Count, tempSelected);
    }

    private async Task ReloadThumbnails()
    {
        LoadingAnnsCaption.Visibility = Visibility.Visible;
        LoadingAnnsBar.Visibility = Visibility.Visible;

        if (!Directory.Exists(_config.ThumbnailsDirectory))
            return;

        var thumbnails = Directory.GetFiles(_config.ThumbnailsDirectory, "*.jpg");
        var thumbnailDtos = new List<ThumbnailDto>();
        for (int i = 0; i < thumbnails.Length; i++)
        {
            var thumbnailDto = await GetThumbnail(thumbnails[i]);
            if (thumbnailDto != null)
                thumbnailDtos.Add(thumbnailDto);

            if (i % 1000 == 0)
                LoadingAnnsBar.Value = i * 100.0 / thumbnails.Length;
        }

        ThumbnailsDtos.Clear();
        foreach (var th in thumbnailDtos.OrderByDescending(x => x.ImageDate))
            ThumbnailsDtos.Add(th);

        LoadingAnnsCaption.Visibility = Visibility.Collapsed;
        LoadingAnnsBar.Visibility = Visibility.Collapsed;
    }

    private async Task<ThumbnailDto?> GetThumbnail(string thumbnail)
    {
        try
        {
            var name = Path.GetFileNameWithoutExtension(thumbnail)[..^Config.THUMBNAIL_PREFIX.Length];
            var imagePath = Path.Combine(_config.ImagesDirectory, name);
            var labelPath = Path.Combine(_config.LabelsDirectory, $"{name}.txt");

            foreach (var f in _config.ImageFormats)
            {
                var curName = $"{imagePath}.{f}";
                if (File.Exists(curName))
                {
                    imagePath = curName;
                    break;
                }
            }

            if (!_galleryManager.LabelsCache.TryGetValue(Path.GetFileName(imagePath), out var info))
            {
                if (!File.Exists(imagePath) || !File.Exists(labelPath))
                {
                    File.Delete(thumbnail);
                    _logger.LogError($"No label {labelPath} found ! Image {imagePath} not found, thumbnail {thumbnail} was removed");
                    return null;
                }

                var classes =  (await YoloLabel.ReadFromFile(labelPath))
                    .Select(x => x.ClassNumber)
                    .Distinct()
                    .ToList();

                info = _galleryManager.AddToCache(imagePath, classes);
            }

            if (!info.Classes.Contains(ExplorerEditor.CurrentAnnClass.Id) && ExplorerEditor.CurrentAnnClass.Id != -1)
                return null;

            return new ThumbnailDto
            {
                ThumbnailPath = thumbnail,
                ImagePath = imagePath,
                LabelPath = labelPath,
                ImageDate = info.ImageDateTime
            };
        }
        catch (Exception e)
        {
            _logger.LogError(e, e.Message);
            return null;
        }
    }

    public void AddThumbnail(ThumbnailDto thumbnailDto, IEnumerable<int> classes)
    {
        var selectedClass = ((AnnotationClass?)LvClasses.SelectedItem)?.Id;

        if (selectedClass != null && (selectedClass == -1 || classes.Any(x => x == selectedClass)))
            ThumbnailsDtos.Insert(0, thumbnailDto);
    }
}