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 Newtonsoft.Json;
using ScottPlot;
using Color = ScottPlot.Color;
using MessageBox = System.Windows.MessageBox;

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 string _thumbnailsCacheFile;
    private readonly IConfigRepository _configRepository;
    private readonly FormState _formState;
    private static Dictionary<string, List<int>> LabelsCache { get; set; } = new();

    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;
        _thumbnailsCacheFile = Path.Combine(config.ThumbnailsDirectory, Config.ThumbnailsCacheFile);
        if (File.Exists(_thumbnailsCacheFile))
        {
            var cache = JsonConvert.DeserializeObject<Dictionary<string, List<int>>>(File.ReadAllText(_thumbnailsCacheFile));
            LabelsCache = cache ?? new Dictionary<string, List<int>>();
        }

        InitializeComponent();
        Loaded += async (_, _) =>
        {
            AllAnnotationClasses = new ObservableCollection<AnnotationClass>(
                new List<AnnotationClass> { new(-1, "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; };

        Switcher.SelectionChanged += (sender, args) =>
        {
            switch (Switcher.SelectedIndex)
            {
                case 0: //ListView
                    LvClasses.ItemsSource = AllAnnotationClasses;
                    LvClasses.SelectedIndex = _tempSelectedClassIdx;
                    ExplorerEditor.Background = null;
                    break;
                case 1: //Editor
                    _tempSelectedClassIdx = LvClasses.SelectedIndex;
                    LvClasses.ItemsSource = _config.AnnotationClasses;
                    break;
            }
        };

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

    }

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

        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
        }));
        foreach (var x in data)
        {
            var label = plot.Add.Text(x.Name, 50, -1.5 * x.Key + 1.1);
            label.LabelFontSize = 16;
        }

        plot.Axes.AutoScale();
        //plot.Axes.SetLimits(-200, data.Max(x => x.ClassCount + 3000), -2 * data.Count + 5, 5);
        ClassDistribution.Background = new SolidColorBrush(System.Windows.Media.Colors.Black);
        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()
            };

            Switcher.SelectedIndex = 1;
            LvClasses.SelectedIndex = 1;

            var time = _formState.GetTime(dto.ImagePath);
            ExplorerEditor.RemoveAllAnns();
            foreach (var ann in await YoloLabel.ReadFromFile(dto.LabelPath))
            {
                var annClass = _config.AnnotationClassesDict[ann.ClassNumber];
                var annInfo = new CanvasLabel(ann, ExplorerEditor.RenderSize, ExplorerEditor.RenderSize);
                ExplorerEditor.CreateAnnotation(annClass, time, annInfo);
            }

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

    }

    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;

        ThumbnailsDtos.Clear();
        var thumbnails = Directory.GetFiles(_config.ThumbnailsDirectory, "*.jpg");

        var thumbNum = 0;
        foreach (var thumbnail in thumbnails)
        {
            await AddThumbnail(thumbnail);

            if (thumbNum % 1000 == 0)
            {
                await File.WriteAllTextAsync(_thumbnailsCacheFile, JsonConvert.SerializeObject(LabelsCache));
                LoadingAnnsBar.Value = thumbNum * 100.0 / thumbnails.Length;
            }
            thumbNum++;
        }

        LoadingAnnsCaption.Visibility = Visibility.Collapsed;
        LoadingAnnsBar.Visibility = Visibility.Collapsed;
        await File.WriteAllTextAsync(_thumbnailsCacheFile, JsonConvert.SerializeObject(LabelsCache));
    }

    private async Task AddThumbnail(string thumbnail, CancellationToken cancellationToken = default)
    {
        try
        {
            var name = Path.GetFileNameWithoutExtension(thumbnail)[..^Config.ThumbnailPrefix.Length];
            var imageName = Path.Combine(_config.ImagesDirectory, name);
            foreach (var f in _config.ImageFormats)
            {
                var curName = $"{imageName}.{f}";
                if (File.Exists(curName))
                {
                    imageName = curName;
                    break;
                }
            }

            var labelPath = Path.Combine(_config.LabelsDirectory, $"{name}.txt");

            if (!LabelsCache.TryGetValue(name, out var classes))
            {
                if (!File.Exists(labelPath))
                {
                    var imageExists = File.Exists(imageName);
                    if (!imageExists)
                    {
                        _logger.LogError($"No label {labelPath} found ! Image {imageName} not found, removing thumbnail {thumbnail}");
                        File.Delete(thumbnail);
                    }
                    else
                    {
                        _logger.LogError($"No label {labelPath} found! But Image {imageName} exists! Image moved to {_config.UnknownImages} directory!");
                        File.Move(imageName, Path.Combine(_config.UnknownImages, imageName));
                    }
                    return;
                }

                var labels = await YoloLabel.ReadFromFile(labelPath, cancellationToken);
                classes = labels.Select(x => x.ClassNumber).Distinct().ToList();
                LabelsCache.Add(name, classes);
            }

            if (classes.Contains(ExplorerEditor.CurrentAnnClass.Id) || ExplorerEditor.CurrentAnnClass.Id == -1)
            {
                ThumbnailsDtos.Add(new ThumbnailDto
                {
                    ThumbnailPath = thumbnail,
                    ImagePath = imageName,
                    LabelPath = labelPath
                });
            }
        }
        catch (Exception e)
        {
            _logger.LogError(e, e.Message);
        }
    }
}