annotations/Azaion.Annotator/MainWindow.xaml.cs

using System.Collections.ObjectModel;
using System.IO;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Controls.Primitives;
using System.Windows.Input;
using Azaion.Annotator.DTO;
using Azaion.Annotator.Extensions;
using LibVLCSharp.Shared;
using MediatR;
using Microsoft.WindowsAPICodePack.Dialogs;
using Newtonsoft.Json;
using Point = System.Windows.Point;
using Size = System.Windows.Size;
using IntervalTree;
using Microsoft.Extensions.Logging;

namespace Azaion.Annotator;

public partial class MainWindow
{
    private readonly LibVLC _libVLC;
    private readonly MediaPlayer _mediaPlayer;
    private readonly IMediator _mediator;
    private readonly FormState _formState;

    private readonly IConfigRepository _configRepository;
    private readonly HelpWindow _helpWindow;
    private readonly ILogger<MainWindow> _logger;
    private readonly IGalleryManager _galleryManager;
    private CancellationTokenSource _cancellationTokenSource = new CancellationTokenSource();

    private ObservableCollection<AnnotationClass> AnnotationClasses { get; set; } = new();
    private bool _suspendLayout;

    private readonly TimeSpan _thresholdBefore = TimeSpan.FromMilliseconds(100);
    private readonly TimeSpan _thresholdAfter = TimeSpan.FromMilliseconds(300);
    private readonly Config _config;
    private readonly DatasetExplorer _datasetExplorer;

    private ObservableCollection<MediaFileInfo> AllMediaFiles { get; set; } = new();
    private ObservableCollection<MediaFileInfo> FilteredMediaFiles { get; set; } = new();

    public IntervalTree<TimeSpan, List<YoloLabel>> Annotations { get; set; } = new();

    public MainWindow(LibVLC libVLC, MediaPlayer mediaPlayer,
        IMediator mediator,
        FormState formState,
        IConfigRepository configRepository,
        HelpWindow helpWindow,
        DatasetExplorer datasetExplorer,
        ILogger<MainWindow> logger,
        IGalleryManager galleryManager)
    {
        InitializeComponent();
        _libVLC = libVLC;
        _mediaPlayer = mediaPlayer;
        _mediator = mediator;
        _formState = formState;
        _configRepository = configRepository;
        _config = _configRepository.Get();
        _helpWindow = helpWindow;
        _datasetExplorer = datasetExplorer;
        _logger = logger;
        _galleryManager = galleryManager;

        VideoView.Loaded += VideoView_Loaded;
        Closed += OnFormClosed;

        if (!Directory.Exists(_config.LabelsDirectory))
            Directory.CreateDirectory(_config.LabelsDirectory);
        if (!Directory.Exists(_config.ImagesDirectory))
            Directory.CreateDirectory(_config.ImagesDirectory);
        if (!Directory.Exists(_config.ResultsDirectory))
            Directory.CreateDirectory(_config.ResultsDirectory);

        Editor.GetTimeFunc = () => TimeSpan.FromMilliseconds(_mediaPlayer.Time);

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

    private void VideoView_Loaded(object sender, RoutedEventArgs e)
    {
        Core.Initialize();
        InitControls();

        _ = Task.Run(async () =>
        {
            while (true)
            {
                await _galleryManager.RefreshThumbnails();
                await Task.Delay(30000);
            }
        });

        _suspendLayout = true;

        Left = _config.MainWindowConfig.WindowLocation.X;
        Top = _config.MainWindowConfig.WindowLocation.Y;
        Width = _config.MainWindowConfig.WindowSize.Width;
        Height = _config.MainWindowConfig.WindowSize.Height;

        _datasetExplorer.Left = _config.MainWindowConfig.WindowLocation.X;
        _datasetExplorer.Top = _config.DatasetExplorerConfig.WindowLocation.Y;
        _datasetExplorer.Width = _config.DatasetExplorerConfig.WindowSize.Width;
        _datasetExplorer.Height = _config.DatasetExplorerConfig.WindowSize.Height;
        if (_config.DatasetExplorerConfig.FullScreen)
            _datasetExplorer.WindowState = WindowState.Maximized;

        MainGrid.ColumnDefinitions.FirstOrDefault()!.Width = new GridLength(_config.LeftPanelWidth);
        MainGrid.ColumnDefinitions.LastOrDefault()!.Width = new GridLength(_config.RightPanelWidth);

        if (_config.MainWindowConfig.FullScreen)
            WindowState = WindowState.Maximized;

        _suspendLayout = false;

        ReloadFiles();
        if (_config.AnnotationClasses.Count == 0)
            _config.AnnotationClasses.Add(new AnnotationClass(0));

        AnnotationClasses = new ObservableCollection<AnnotationClass>(_config.AnnotationClasses);
        LvClasses.ItemsSource = AnnotationClasses;
        LvClasses.SelectedIndex = 0;

        if (LvFiles.Items.IsEmpty) 
            BlinkHelp(HelpTexts.HelpTextsDict[HelpTextEnum.Initial]);
        
        if (_config.ShowHelpOnStart) 
            _helpWindow.Show();
    }

    public void BlinkHelp(string helpText, int times = 2)
    {
        _ = Task.Run(async () =>
        {
            for (int i = 0; i < times; i++)
            {
                Dispatcher.Invoke(() => StatusHelp.Text = helpText);
                await Task.Delay(200);
                Dispatcher.Invoke(() => StatusHelp.Text = "");
                await Task.Delay(200);
            }
            Dispatcher.Invoke(() => StatusHelp.Text = helpText);
        });
    }

    private void InitControls()
    {
        VideoView.MediaPlayer = _mediaPlayer;

        _mediaPlayer.Playing += async (sender, args) =>
        {
            if (_formState.CurrentMrl == _mediaPlayer.Media?.Mrl)
                return; //already loaded all the info

            _formState.CurrentMrl = _mediaPlayer.Media?.Mrl ?? "";
            uint vw = 0, vh = 0;
            _mediaPlayer.Size(0, ref vw, ref vh);
            _formState.CurrentVideoSize = new Size(vw, vh);
            _formState.CurrentVideoLength = TimeSpan.FromMilliseconds(_mediaPlayer.Length);

            await Dispatcher.Invoke(async () => await ReloadAnnotations(_cancellationTokenSource.Token));
            if (_formState.CurrentMedia?.MediaType != MediaTypes.Image)
                return;

            //if image show annotations, give 100ms to load the frame and set on pause
            await Task.Delay(100);
            ShowCurrentAnnotations();
            _mediaPlayer.SetPause(true);
        };

        LvFiles.MouseDoubleClick += async (_, _) => await _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.Play));
        
        LvClasses.SelectionChanged += (_, _) =>
        {
            var selectedClass = (AnnotationClass)LvClasses.SelectedItem;
            Editor.CurrentAnnClass = selectedClass;
            _mediator.Publish(new AnnClassSelectedEvent(selectedClass));
        };

        _mediaPlayer.PositionChanged += (o, args) =>
        {
            ShowTimeAnnotations(TimeSpan.FromMilliseconds(_mediaPlayer.Time));
        };
        
        VideoSlider.ValueChanged += (value, newValue) => 
            _mediaPlayer.Position = (float)(newValue / VideoSlider.Maximum);

        VideoSlider.KeyDown += (sender, args) => _mediator.Publish(new KeyEvent(sender, args));

        Volume.ValueChanged += (_, newValue) => _mediator.Publish(new VolumeChangedEvent((int)newValue));

        SizeChanged += async (_, _) => await SaveUserSettings();
        LocationChanged += async (_, _) => await SaveUserSettings();
        StateChanged += async (_, _) => await SaveUserSettings();
        
        Editor.FormState = _formState;
        Editor.Mediator = _mediator;
        DgAnnotations.ItemsSource = _formState.AnnotationResults;
    }

    private async Task SaveUserSettings()
    {
        if (_suspendLayout)
            return;

        _config.LeftPanelWidth = MainGrid.ColumnDefinitions.FirstOrDefault()!.Width.Value;
        _config.RightPanelWidth = MainGrid.ColumnDefinitions.LastOrDefault()!.Width.Value;

        _config.MainWindowConfig = this.GetConfig();
        await ThrottleExt.Throttle(() =>
        {
            _configRepository.Save(_config);
            return Task.CompletedTask;
        }, TimeSpan.FromSeconds(5));
    }

    public void ShowCurrentAnnotations() => ShowTimeAnnotations(TimeSpan.FromMilliseconds(_mediaPlayer.Time));

    private void ShowTimeAnnotations(TimeSpan time)
    {
        Dispatcher.Invoke(() => VideoSlider.Value = _mediaPlayer.Position * VideoSlider.Maximum);
        Dispatcher.Invoke(() => StatusClock.Text = $"{TimeSpan.FromMilliseconds(_mediaPlayer.Time):mm\\:ss} / {_formState.CurrentVideoLength:mm\\:ss}");

        Dispatcher.Invoke(() => Editor.ClearExpiredAnnotations(time));

        var annotations = Annotations.Query(time).SelectMany(x => x).ToList();
        foreach (var ann in annotations)
        {
            var annClass = _config.AnnotationClasses[ann.ClassNumber];
            var annInfo = new CanvasLabel(ann, Editor.RenderSize, _formState.CurrentVideoSize);
            Dispatcher.Invoke(() => Editor.CreateAnnotation(annClass, time, annInfo));
        }
    }

    public async Task ReloadAnnotations(CancellationToken cancellationToken)
    {
        _formState.AnnotationResults.Clear();
        Annotations.Clear();
        Editor.RemoveAllAnns();

        var labelDir = new DirectoryInfo(_config.LabelsDirectory);
        if (!labelDir.Exists)
            return;
        
        var labelFiles = labelDir.GetFiles($"{_formState.VideoName}_??????.txt");
        foreach (var file in labelFiles)
        {
            var name = Path.GetFileNameWithoutExtension(file.Name);
            var time = _formState.GetTime(name)!.Value;

            await AddAnnotation(time, await YoloLabel.ReadFromFile(file.FullName));
        }
    }

    public async Task AddAnnotation(TimeSpan time, List<YoloLabel> annotations)
    {
        var fName = _formState.GetTimeName(time);

        var previousAnnotations = Annotations.Query(time);
        Annotations.Remove(previousAnnotations);
        Annotations.Add(time.Subtract(_thresholdBefore), time.Add(_thresholdAfter), annotations);

        var existingResult = _formState.AnnotationResults.FirstOrDefault(x => x.Time == time);
        if (existingResult != null)
            _formState.AnnotationResults.Remove(existingResult);

        _formState.AnnotationResults.Add(new AnnotationResult(time, fName, annotations, _config));
        await File.WriteAllTextAsync($"{_config.ResultsDirectory}/{fName}.json", JsonConvert.SerializeObject(_formState.AnnotationResults));
    }

    private void ReloadFiles()
    {
        var dir = new DirectoryInfo(_config.VideosDirectory);
        if (!dir.Exists)
            return;

        var labelNames = new DirectoryInfo(_config.LabelsDirectory).GetFiles()
            .Select(x => x.Name[..^11])
            .GroupBy(x => x)
            .Select(gr => gr.Key)
            .ToDictionary(x => x);

        var videoFiles = dir.GetFiles(_config.VideoFormats.ToArray()).Select(x =>
        {
            var media = new Media(_libVLC, x.FullName);
            media.Parse();
            var fInfo = new MediaFileInfo
            {
                Name = x.Name,
                Path = x.FullName,
                MediaType = MediaTypes.Video,
                HasAnnotations = labelNames.ContainsKey(Path.GetFileNameWithoutExtension(x.Name).Replace(" ", ""))
            };
            media.ParsedChanged += (_, _) => fInfo.Duration = TimeSpan.FromMilliseconds(media.Duration);
            return fInfo;
        }).ToList();

        var imageFiles = dir.GetFiles(_config.ImageFormats.ToArray()).Select(x => new MediaFileInfo
        {
            Name = x.Name,
            Path = x.FullName,
            MediaType = MediaTypes.Image,
            HasAnnotations = labelNames.ContainsKey(Path.GetFileNameWithoutExtension(x.Name).Replace(" ", ""))
        });

        AllMediaFiles = new ObservableCollection<MediaFileInfo>(videoFiles.Concat(imageFiles).ToList());
        LvFiles.ItemsSource = AllMediaFiles;
        TbFolder.Text = _config.VideosDirectory;

        BlinkHelp(AllMediaFiles.Count == 0
            ? HelpTexts.HelpTextsDict[HelpTextEnum.Initial] 
            : HelpTexts.HelpTextsDict[HelpTextEnum.PlayVideo]);
    }

    private void OnFormClosed(object? sender, EventArgs e)
    {
        _mediaPlayer.Stop();
        _mediaPlayer.Dispose();
        _libVLC.Dispose();
        _configRepository.Save(_config);
        Application.Current.Shutdown();
    }

    // private void AddClassBtnClick(object sender, RoutedEventArgs e)
    // {
    //     LvClasses.IsReadOnly = false;
    //     AnnotationClasses.Add(new AnnotationClass(AnnotationClasses.Count));
    //     LvClasses.SelectedIndex = AnnotationClasses.Count - 1;
    // }
    private async void OpenFolderItemClick(object sender, RoutedEventArgs e) => await OpenFolder();
    private async void OpenFolderButtonClick(object sender, RoutedEventArgs e) => await OpenFolder();
    private async Task OpenFolder()
    {
        var dlg = new CommonOpenFileDialog
        {
            Title = "Open Video folder",
            IsFolderPicker = true,
            InitialDirectory = Path.GetDirectoryName(AppDomain.CurrentDomain.BaseDirectory)
        };
        if (dlg.ShowDialog() != CommonFileDialogResult.Ok)
            return;

        if (!string.IsNullOrEmpty(dlg.FileName))
        {
            _config.VideosDirectory = dlg.FileName;
            await SaveUserSettings();
        }

        ReloadFiles();
    }
    private void TbFilter_OnTextChanged(object sender, TextChangedEventArgs e)
    {
        FilteredMediaFiles = new ObservableCollection<MediaFileInfo>(AllMediaFiles.Where(x => x.Name.ToLower().Contains(TbFilter.Text.ToLower())).ToList());
        LvFiles.ItemsSource = FilteredMediaFiles;
    }

    private void OpenDataExplorerItemClick(object sender, RoutedEventArgs e)
    {
        _datasetExplorer.Show();
        _datasetExplorer.Activate();
    }


    private void PlayClick(object sender, RoutedEventArgs e)
    {
        _mediator.Publish(new PlaybackControlEvent(_mediaPlayer.CanPause ? PlaybackControlEnum.Pause : PlaybackControlEnum.Play));
    }

    private void PauseClick(object sender, RoutedEventArgs e) => _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.Pause));
    private void StopClick(object sender, RoutedEventArgs e) => _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.Stop));

    private void PreviousFrameClick(object sender, RoutedEventArgs e) => _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.PreviousFrame));
    private void NextFrameClick(object sender, RoutedEventArgs e) =>  _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.NextFrame));

    private void SaveAnnotationsClick(object sender, RoutedEventArgs e) => _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.SaveAnnotations));

    private void RemoveSelectedClick(object sender, RoutedEventArgs e) => _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.RemoveSelectedAnns));
    private void RemoveAllClick(object sender, RoutedEventArgs e) => _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.RemoveAllAnns));
    private void TurnOffVolume(object sender, RoutedEventArgs e) =>  _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.TurnOffVolume));
    private void TurnOnVolume(object sender, RoutedEventArgs e) =>  _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.TurnOnVolume));
    
    private void OpenHelpWindowClick(object sender, RoutedEventArgs e)
    {
        _helpWindow.Show();
        _helpWindow.Activate();
    }

    private void DgAnnotationsRowClick(object sender, MouseButtonEventArgs e)
    {
        DgAnnotations.MouseDoubleClick += (sender, args) =>
        {
            Editor.RemoveAllAnns();
            var dgRow = ItemsControl.ContainerFromElement((DataGrid)sender, (args.OriginalSource as DependencyObject)!) as DataGridRow;
            var res = (AnnotationResult)dgRow!.Item;
            _mediaPlayer.SetPause(true);
            _mediaPlayer.Time = (long)res.Time.TotalMilliseconds;// + 250;
            ShowTimeAnnotations(res.Time);
        };
    }

    private void Thumb_OnDragCompleted(object sender, DragCompletedEventArgs e) => _ = SaveUserSettings();
}