annotations/Azaion.Annotator/Annotator.xaml.cs

using System.Collections.ObjectModel;
using System.Diagnostics;
using System.Drawing.Imaging;
using System.IO;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Controls.Primitives;
using System.Windows.Input;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Windows.Threading;
using Azaion.Annotator.DTO;
using Azaion.Annotator.Extensions;
using Azaion.Common;
using Azaion.Common.DTO;
using Azaion.Common.DTO.Config;
using Azaion.Common.Extensions;
using LibVLCSharp.Shared;
using MediatR;
using Microsoft.WindowsAPICodePack.Dialogs;
using Newtonsoft.Json;
using Size = System.Windows.Size;
using IntervalTree;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using MediaPlayer = LibVLCSharp.Shared.MediaPlayer;

namespace Azaion.Annotator;

public partial class Annotator
{
    private readonly AppConfig _appConfig;
    private readonly LibVLC _libVLC;
    private readonly MediaPlayer _mediaPlayer;
    private readonly IMediator _mediator;
    private readonly FormState _formState;

    private readonly IConfigUpdater _configUpdater;
    private readonly HelpWindow _helpWindow;
    private readonly ILogger<Annotator> _logger;
    private readonly VLCFrameExtractor _vlcFrameExtractor;
    private readonly IAIDetector _aiDetector;
    private readonly CancellationTokenSource _cancellationTokenSource = new();

    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 ObservableCollection<MediaFileInfo> AllMediaFiles { get; set; } = new();
    private ObservableCollection<MediaFileInfo> FilteredMediaFiles { get; set; } = new();

    public IntervalTree<TimeSpan, List<YoloLabel>> Annotations { get; set; } = new();
    private AutodetectDialog _autoDetectDialog = new() { Topmost = true };

    public Annotator(
        IConfigUpdater configUpdater,
        IOptions<AppConfig> appConfig,
        LibVLC libVLC, MediaPlayer mediaPlayer,
        IMediator mediator,
        FormState formState,
        HelpWindow helpWindow,
        ILogger<Annotator> logger,
        VLCFrameExtractor vlcFrameExtractor,
        IAIDetector aiDetector)
    {
        InitializeComponent();
        _appConfig = appConfig.Value;
        _configUpdater = configUpdater;
        _libVLC = libVLC;
        _mediaPlayer = mediaPlayer;
        _mediator = mediator;
        _formState = formState;
        _helpWindow = helpWindow;
        _logger = logger;
        _vlcFrameExtractor = vlcFrameExtractor;
        _aiDetector = aiDetector;

        Loaded += OnLoaded;
        Closed += OnFormClosed;
        Activated += (_, _) => _formState.ActiveWindow = WindowEnum.Annotator;
        Editor.GetTimeFunc = () => TimeSpan.FromMilliseconds(_mediaPlayer.Time);
    }

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

        _suspendLayout = true;

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

        _suspendLayout = false;

        ReloadFiles();

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

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

    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;

        //On start playing media
        _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)
            {
                await Task.Delay(100); //wait to load the frame and set on pause
                ShowTimeAnnotations(TimeSpan.FromMilliseconds(_mediaPlayer.Time));
                _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, WindowEnum.Annotator));

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

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

        DgAnnotations.MouseDoubleClick += (sender, args) =>
        {
            var dgRow = ItemsControl.ContainerFromElement((DataGrid)sender, (args.OriginalSource as DependencyObject)!) as DataGridRow;
            OpenAnnotationResult((AnnotationResult)dgRow!.Item);
        };

        DgAnnotations.KeyUp += (sender, args) =>
        {
            switch (args.Key)
            {
                case Key.Up:
                case Key.Down: //cursor is already moved by system behaviour
                    OpenAnnotationResult((AnnotationResult)DgAnnotations.SelectedItem);
                    break;
                case Key.Delete:
                    var result = MessageBox.Show(Application.Current.MainWindow, "Чи дійсно видалити аннотації?","Підтвердження видалення", MessageBoxButton.OKCancel, MessageBoxImage.Question);
                    if (result != MessageBoxResult.OK)
                        return;

                    // var allWindows = Application.Current.Windows.Cast<Window>();
                    // try
                    // {
                    //     foreach (var window in allWindows)
                    //         window.IsEnabled = false;
                    //
                    // }
                    // finally
                    // {
                    //     foreach (var window in allWindows)
                    //     {
                    //         window.IsEnabled = true;
                    //     }
                    // }

                    var res = DgAnnotations.SelectedItems.Cast<AnnotationResult>().ToList();
                    foreach (var annotationResult in res)
                    {
                        var imgName = Path.GetFileNameWithoutExtension(annotationResult.Image);
                        var thumbnailPath = Path.Combine(_appConfig.DirectoriesConfig.ThumbnailsDirectory, $"{imgName}{Constants.THUMBNAIL_PREFIX}.jpg");
                        File.Delete(annotationResult.Image);
                        File.Delete(Path.Combine(_appConfig.DirectoriesConfig.LabelsDirectory, $"{imgName}.txt"));
                        File.Delete(thumbnailPath);
                        _formState.AnnotationResults.Remove(annotationResult);
                        Annotations.Remove(Annotations.Query(annotationResult.Time));
                    }
                    break;
            }
        };

        Editor.Mediator = _mediator;
        DgAnnotations.ItemsSource = _formState.AnnotationResults;
    }

    public void OpenAnnotationResult(AnnotationResult res)
    {
        _mediaPlayer.SetPause(true);
        Editor.RemoveAllAnns();
        _mediaPlayer.Time = (long)res.Time.TotalMilliseconds;

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

        AddAnnotationsToCanvas(res.Time, res.Detections, showImage: true);
    }
    private async Task SaveUserSettings()
    {
        if (_suspendLayout)
            return;

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

        await ThrottleExt.Throttle(() =>
        {
            _configUpdater.Save(_appConfig);
            return Task.CompletedTask;
        }, TimeSpan.FromSeconds(5));
    }

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

        var annotations = Annotations.Query(time).SelectMany(x => x).Select(x => new Detection(x));
        AddAnnotationsToCanvas(time, annotations);
    }

    private void AddAnnotationsToCanvas(TimeSpan? time, IEnumerable<Detection> labels, bool showImage = false)
    {
        Dispatcher.Invoke(async () =>
        {
            var canvasSize = Editor.RenderSize;
            var videoSize = _formState.CurrentVideoSize;
            if (showImage)
            {
                var fName = _formState.GetTimeName(time);
                var imgPath = Path.Combine(_appConfig.DirectoriesConfig.ImagesDirectory, $"{fName}.jpg");
                if (File.Exists(imgPath))
                {
                    Editor.Background = new ImageBrush { ImageSource = await imgPath.OpenImage() };
                    _formState.BackgroundTime = time;
                    videoSize = Editor.RenderSize;
                }
            }
            foreach (var label in labels)
            {
                var annClass = _appConfig.AnnotationConfig.AnnotationClasses[label.ClassNumber];
                var canvasLabel = new CanvasLabel(label, canvasSize, videoSize, label.Probability);
                Editor.CreateAnnotation(annClass, time, canvasLabel);
            }

        });
    }

    private async Task ReloadAnnotations(CancellationToken ct = default)
    {
        _formState.AnnotationResults.Clear();
        Annotations.Clear();
        Editor.RemoveAllAnns();

        var labelDir = new DirectoryInfo(_appConfig.DirectoriesConfig.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 = Constants.GetTime(name);
            await AddAnnotations(time, await YoloLabel.ReadFromFile(file.FullName, ct), ct);
        }
    }

    public async Task AddAnnotations(TimeSpan? time, List<YoloLabel> annotations, CancellationToken ct = default)
        => await AddAnnotations(time, annotations.Select(x => new Detection(x)).ToList(), ct);

    public async Task AddAnnotations(TimeSpan? time, List<Detection> detections, CancellationToken ct = default)
    {
        var timeValue = time ?? TimeSpan.FromMinutes(0);
        var previousAnnotations = Annotations.Query(timeValue);
        Annotations.Remove(previousAnnotations);
        Annotations.Add(timeValue.Subtract(_thresholdBefore), timeValue.Add(_thresholdAfter), detections.Cast<YoloLabel>().ToList());

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

        var dict = _formState.AnnotationResults
            .Select((x, i) => new { x.Time, Index = i })
            .ToDictionary(x => x.Time, x => x.Index);

        var index = dict.Where(x => x.Key < timeValue)
            .OrderBy(x => timeValue - x.Key)
            .Select(x => x.Value + 1)
            .FirstOrDefault();

        _formState.AnnotationResults.Insert(index, CreateAnnotationReult(timeValue, detections));
        await File.WriteAllTextAsync($"{_appConfig.DirectoriesConfig.ResultsDirectory}/{_formState.VideoName}.json", JsonConvert.SerializeObject(_formState.AnnotationResults), ct);
    }

    private AnnotationResult CreateAnnotationReult(TimeSpan timeValue, List<Detection> detections)
    {
        var annotationResult = new AnnotationResult
        {
            Time = timeValue,
            Image = $"{_formState.GetTimeName(timeValue)}.jpg",
            Detections = detections,
        };
        if (detections.Count <= 0)
            return annotationResult;

        Color GetAnnotationClass(List<int> detectionClasses, int colorNumber)
        {
            if (detections.Count == 0)
                return (-1).ToColor();

            return colorNumber >= detectionClasses.Count
                ? _appConfig.AnnotationConfig.AnnotationClassesDict[detectionClasses.LastOrDefault()].Color
                : _appConfig.AnnotationConfig.AnnotationClassesDict[detectionClasses[colorNumber]].Color;
        }

        var detectionClasses = detections.Select(x => x.ClassNumber).Distinct().ToList();

        annotationResult.ClassName = detectionClasses.Count > 1
                ? string.Join(", ", detectionClasses.Select(x => _appConfig.AnnotationConfig.AnnotationClassesDict[x].ShortName))
                : _appConfig.AnnotationConfig.AnnotationClassesDict[detectionClasses.FirstOrDefault()].Name;

        annotationResult.ClassColor0 = GetAnnotationClass(detectionClasses, 0);
        annotationResult.ClassColor1 = GetAnnotationClass(detectionClasses, 1);
        annotationResult.ClassColor2 = GetAnnotationClass(detectionClasses, 2);
        annotationResult.ClassColor3 = GetAnnotationClass(detectionClasses, 3);
        return annotationResult;
    }

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

        var labelNames = new DirectoryInfo(_appConfig.DirectoriesConfig.LabelsDirectory).GetFiles()
            .Select(x =>
            {
                var name = Path.GetFileNameWithoutExtension(x.Name);
                return name.Length > 8
                    ? name[..^7]
                    : name;
            })
            .GroupBy(x => x)
            .Select(gr => gr.Key)
            .ToDictionary(x => x);

        var videoFiles = dir.GetFiles(_appConfig.AnnotationConfig.VideoFormats.ToArray()).Select(x =>
        {
            using 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(_appConfig.AnnotationConfig.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 = _appConfig.DirectoriesConfig.VideosDirectory;

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

    private void OnFormClosed(object? sender, EventArgs e)
    {
        _mediaPlayer.Stop();
        _mediaPlayer.Dispose();
        _libVLC.Dispose();
    }

    private void OpenContainingFolder(object sender, RoutedEventArgs e)
    {
        var mediaFileInfo = (sender as MenuItem)?.DataContext as MediaFileInfo;
        if (mediaFileInfo == null)
            return;

        Process.Start("explorer.exe", "/select,\"" + mediaFileInfo.Path +"\"");
    }

    public void SeekTo(long timeMilliseconds)
    {
        _mediaPlayer.SetPause(true);
        _mediaPlayer.Time = timeMilliseconds;
        VideoSlider.Value = _mediaPlayer.Position * 100;
    }

    private void SeekTo(TimeSpan time) =>
        SeekTo((long)time.TotalMilliseconds);

    // 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))
        {
            _appConfig.DirectoriesConfig.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 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 Thumb_OnDragCompleted(object sender, DragCompletedEventArgs e) => _ = SaveUserSettings();

    private void LvFilesContextOpening(object sender, ContextMenuEventArgs e)
    {
        var listItem = sender as ListViewItem;
        LvFilesContextMenu.DataContext = listItem!.DataContext;
    }

    private (TimeSpan Time, List<Detection> Detections)? _previousDetection;

    public async void AutoDetect(object sender, RoutedEventArgs e)
    {
        if (LvFiles.Items.IsEmpty)
            return;
        if (LvFiles.SelectedIndex == -1)
            LvFiles.SelectedIndex = 0;

        await _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.Play));
        _mediaPlayer.SetPause(true);

        var manualCancellationSource = new CancellationTokenSource();
        var token = manualCancellationSource.Token;

        _autoDetectDialog = new AutodetectDialog
        {
            Topmost = true,
            Owner = this
        };
        _autoDetectDialog.Closing += (_, _) =>
        {
            manualCancellationSource.Cancel();
            _mediaPlayer.SeekTo(TimeSpan.Zero);
            Editor.RemoveAllAnns();
        };

        _autoDetectDialog.Top = Height - _autoDetectDialog.Height - 80;
        _autoDetectDialog.Left = 5;

        _autoDetectDialog.Log("Ініціалізація AI...");

        _ = Task.Run(async () =>
        {
            var mediaInfo = Dispatcher.Invoke(() => (MediaFileInfo)LvFiles.SelectedItem);
            while (mediaInfo != null)
            {
                _formState.CurrentMedia = mediaInfo;
                await Dispatcher.Invoke(async () => await ReloadAnnotations(token));

                if (mediaInfo.MediaType == MediaTypes.Image)
                {
                    await DetectImage(mediaInfo, manualCancellationSource, token);
                    await Task.Delay(70, token);
                }
                else
                    await DetectVideo(mediaInfo, manualCancellationSource, token);

                mediaInfo = Dispatcher.Invoke(() =>
                {
                    if (LvFiles.SelectedIndex == LvFiles.Items.Count - 1)
                        return null;
                    LvFiles.SelectedIndex += 1;
                    return (MediaFileInfo)LvFiles.SelectedItem;
                });
            }
            Dispatcher.Invoke(() =>
            {
                _autoDetectDialog.Close();
                _mediaPlayer.Stop();
                LvFiles.Items.Refresh();
            });
        }, token);

        _autoDetectDialog.ShowDialog();
        Dispatcher.Invoke(() => Editor.ResetBackground());
    }

    private async Task DetectImage(MediaFileInfo mediaInfo, CancellationTokenSource manualCancellationSource, CancellationToken token)
    {
        try
        {
            var stream = new FileStream(mediaInfo.Path, FileMode.Open);
            var detections = await _aiDetector.Detect(stream, token);
            await ProcessDetection((TimeSpan.FromMilliseconds(0), stream), detections, token);
            if (detections.Count != 0)
                mediaInfo.HasAnnotations = true;
        }
        catch (Exception e)
        {
            _logger.LogError(e, e.Message);
            await manualCancellationSource.CancelAsync();
        }
    }

    private async Task DetectVideo(MediaFileInfo mediaInfo, CancellationTokenSource manualCancellationSource, CancellationToken token)
    {
        var prevSeekTime = 0.0;
        await foreach (var timeframe in _vlcFrameExtractor.ExtractFrames(mediaInfo.Path, token))
        {
            Console.WriteLine($"Detect time: {timeframe.Time}");
            try
            {
                var detections = await _aiDetector.Detect(timeframe.Stream, token);
                var isValid = IsValidDetection(timeframe.Time, detections);

                if (timeframe.Time.TotalSeconds > prevSeekTime + 1)
                {
                    Dispatcher.Invoke(() => SeekTo(timeframe.Time));
                    prevSeekTime = timeframe.Time.TotalSeconds;
                    if (!isValid) //Show frame anyway
                    {
                        var bitmap = new BitmapImage();
                        bitmap.BeginInit();
                        timeframe.Stream.Seek(0, SeekOrigin.Begin);
                        bitmap.StreamSource = timeframe.Stream;
                        bitmap.CacheOption = BitmapCacheOption.OnLoad;
                        bitmap.EndInit();
                        bitmap.Freeze();

                        Dispatcher.Invoke(() =>
                        {
                            Editor.RemoveAllAnns();
                            Editor.Background = new ImageBrush { ImageSource = bitmap };
                        });
                    }
                }

                if (!isValid)
                    continue;

                mediaInfo.HasAnnotations = true;
                await ProcessDetection(timeframe, detections, token);
            }
            catch (Exception ex)
            {
                _logger.LogError(ex, ex.Message);
                await manualCancellationSource.CancelAsync();
            }
        }
    }

    private bool IsValidDetection(TimeSpan time, List<Detection> detections)
    {
        // No AI detection, forbid
        if (detections.Count == 0)
            return false;

        // Very first detection, allow
        if (!_previousDetection.HasValue)
            return true;

        var prev = _previousDetection.Value;

        // Time between detections is >= than Frame Recognition Seconds, allow
        if (time >= prev.Time.Add(TimeSpan.FromSeconds(_appConfig.AIRecognitionConfig.FrameRecognitionSeconds)))
            return true;

        // Detection is earlier than previous + FrameRecognitionSeconds.
        // Look to the detections more in detail

        // More detected objects, allow
        if (detections.Count > prev.Detections.Count)
            return true;

        foreach (var det in detections)
        {
            var point = new Point(det.CenterX, det.CenterY);
            var closestObject = prev.Detections
                .Select(p => new
                {
                    Point = p,
                    Distance = point.SqrDistance(new Point(p.CenterX, p.CenterY))
                })
                .OrderBy(x => x.Distance)
                .First();

            // Closest object is farther than Tracking distance confidence, hence it's a different object, allow
            if (closestObject.Distance > _appConfig.AIRecognitionConfig.TrackingDistanceConfidence)
                return true;

            // Since closest object within distance confidence, then it is tracking of the same object. Then if recognition probability for the object > increase from previous
            if (det.Probability >= closestObject.Point.Probability + _appConfig.AIRecognitionConfig.TrackingProbabilityIncrease)
                return true;
        }

        return false;
    }

    private async Task ProcessDetection((TimeSpan Time, Stream Stream) timeframe, List<Detection> detections, CancellationToken token = default)
    {
        _previousDetection = (timeframe.Time, detections);
        await Dispatcher.Invoke(async () =>
        {
            try
            {
                var time = timeframe.Time;

                var fName = _formState.GetTimeName(timeframe.Time);
                var imgPath = Path.Combine(_appConfig.DirectoriesConfig.ImagesDirectory, $"{fName}.jpg");
                var img = System.Drawing.Image.FromStream(timeframe.Stream);
                img.Save(imgPath, ImageFormat.Jpeg);
                await YoloLabel.WriteToFile(detections, Path.Combine(_appConfig.DirectoriesConfig.LabelsDirectory, $"{fName}.txt"), token);

                Editor.Background = new ImageBrush { ImageSource = await imgPath.OpenImage() };
                Editor.RemoveAllAnns();
                AddAnnotationsToCanvas(time, detections, true);
                await AddAnnotations(timeframe.Time, detections, token);

                var log = string.Join(Environment.NewLine, detections.Select(det =>
                    $"{_appConfig.AnnotationConfig.AnnotationClassesDict[det.ClassNumber].Name}: " +
                    $"xy=({det.CenterX:F2},{det.CenterY:F2}), " +
                    $"size=({det.Width:F2}, {det.Height:F2}), " +
                    $"prob: {det.Probability:F1}%"));

                Dispatcher.Invoke(() => _autoDetectDialog.Log(log));

                await _mediator.Publish(new ImageCreatedEvent(imgPath), token);
            }
            catch (Exception e)
            {
                _logger.LogError(e, e.Message);
            }
        });
    }
}