add autodetection

2026-04-23 03:26:31 +00:00 · 2024-11-02 13:09:00 +02:00
parent b5b77d9492
commit 418a2116b7
19 changed files with 545 additions and 268 deletions
@@ -1,20 +1,25 @@
 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.Threading;
 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;
+using OpenTK.Graphics.OpenGL;
+using Serilog;
+using MediaPlayer = LibVLCSharp.Shared.MediaPlayer;

 namespace Azaion.Annotator;

@@ -29,7 +34,9 @@ public partial class MainWindow
    private readonly HelpWindow _helpWindow;
    private readonly ILogger<MainWindow> _logger;
    private readonly IGalleryManager _galleryManager;
-    private CancellationTokenSource _cancellationTokenSource = new CancellationTokenSource();
+    private readonly VLCFrameExtractor _vlcFrameExtractor;
+    private readonly IAIDetector _aiDetector;
+    private CancellationTokenSource _cancellationTokenSource = new();

    private ObservableCollection<AnnotationClass> AnnotationClasses { get; set; } = new();
    private bool _suspendLayout;
@@ -43,6 +50,7 @@ public partial class MainWindow
    private ObservableCollection<MediaFileInfo> FilteredMediaFiles { get; set; } = new();

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

    public MainWindow(LibVLC libVLC, MediaPlayer mediaPlayer,
        IMediator mediator,
@@ -51,7 +59,9 @@ public partial class MainWindow
        HelpWindow helpWindow,
        DatasetExplorer datasetExplorer,
        ILogger<MainWindow> logger,
-        IGalleryManager galleryManager)
+        IGalleryManager galleryManager,
+        VLCFrameExtractor vlcFrameExtractor,
+        IAIDetector aiDetector)
    {
        InitializeComponent();
        _libVLC = libVLC;
@@ -64,6 +74,8 @@ public partial class MainWindow
        _datasetExplorer = datasetExplorer;
        _logger = logger;
        _galleryManager = galleryManager;
+        _vlcFrameExtractor = vlcFrameExtractor;
+        _aiDetector = aiDetector;

        VideoView.Loaded += VideoView_Loaded;
        Closed += OnFormClosed;
@@ -189,6 +201,39 @@ public partial class MainWindow
        LocationChanged += async (_, _) => await SaveUserSettings();
        StateChanged += async (_, _) => await SaveUserSettings();

+        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);
+            Editor.RemoveAllAnns();
+            _mediaPlayer.Time = (long)res.Time.TotalMilliseconds;
+            ShowTimeAnnotations(res.Time);
+        };
+
+        DgAnnotations.KeyUp += (sender, args) =>
+        {
+            if (args.Key != Key.Delete)
+                return;
+
+            var result = MessageBox.Show("Чи дійсно видалити аннотації?","Підтвердження видалення", MessageBoxButton.OKCancel, MessageBoxImage.Question);
+            if (result != MessageBoxResult.OK)
+                return;
+
+            var res = DgAnnotations.SelectedItems.Cast<AnnotationResult>().ToList();
+            foreach (var annotationResult in res)
+            {
+                var imgName = Path.GetFileNameWithoutExtension(annotationResult.Image);
+                var thumbnailPath = Path.Combine(_config.ThumbnailsDirectory, $"{imgName}{Config.THUMBNAIL_PREFIX}.jpg");
+                File.Delete(annotationResult.Image);
+                File.Delete(Path.Combine(_config.LabelsDirectory, $"{imgName}.txt"));
+                File.Delete(thumbnailPath);
+                _formState.AnnotationResults.Remove(annotationResult);
+                Annotations.Remove(Annotations.Query(annotationResult.Time));
+            }
+        };
+
        Editor.FormState = _formState;
        Editor.Mediator = _mediator;
        DgAnnotations.ItemsSource = _formState.AnnotationResults;
@@ -221,14 +266,16 @@ public partial class MainWindow

        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));
-        }
+            AddAnnotationToCanvas(time, new CanvasLabel(ann, Editor.RenderSize, _formState.CurrentVideoSize));
    }

-    public async Task ReloadAnnotations(CancellationToken cancellationToken)
+    private void AddAnnotationToCanvas(TimeSpan? time, CanvasLabel canvasLabel)
+    {
+        var annClass = _config.AnnotationClasses[canvasLabel.ClassNumber];
+        Dispatcher.Invoke(() => Editor.CreateAnnotation(annClass, time, canvasLabel));
+    }
+
+    private async Task ReloadAnnotations(CancellationToken cancellationToken)
    {
        _formState.AnnotationResults.Clear();
        Annotations.Clear();
@@ -243,14 +290,12 @@ public partial class MainWindow
        {
            var name = Path.GetFileNameWithoutExtension(file.Name);
            var time = _formState.GetTime(name);
-            await AddAnnotation(time, await YoloLabel.ReadFromFile(file.FullName, cancellationToken));
+            await AddAnnotations(time, await YoloLabel.ReadFromFile(file.FullName, cancellationToken));
        }
    }

-    public async Task AddAnnotation(TimeSpan? time, List<YoloLabel> annotations)
+    public async Task AddAnnotations(TimeSpan? time, List<YoloLabel> annotations)
    {
-        var fName = _formState.GetTimeName(time);
-
        var timeValue = time ?? TimeSpan.FromMinutes(0);
        var previousAnnotations = Annotations.Query(timeValue);
        Annotations.Remove(previousAnnotations);
@@ -269,8 +314,8 @@ public partial class MainWindow
            .Select(x => x.Value + 1)
            .FirstOrDefault();

-        _formState.AnnotationResults.Insert(index, new AnnotationResult(timeValue, fName, annotations, _config));
-        await File.WriteAllTextAsync($"{_config.ResultsDirectory}/{fName}.json", JsonConvert.SerializeObject(_formState.AnnotationResults));
+        _formState.AnnotationResults.Insert(index, new AnnotationResult(timeValue, _formState.GetTimeName(time), annotations, _config));
+        await File.WriteAllTextAsync($"{_config.ResultsDirectory}/{_formState.VideoName}.json", JsonConvert.SerializeObject(_formState.AnnotationResults));
    }

    private void ReloadFiles()
@@ -339,9 +384,19 @@ public partial class MainWindow
        if (mediaFileInfo == null)
            return;

-        Process.Start("explorer.exe", "/select, \"" + mediaFileInfo.Path +"\"");
+        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;
@@ -402,32 +457,12 @@ public partial class MainWindow
    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 async void AutoDetect(object sender, RoutedEventArgs e)
-    {
-        if (LvFiles.SelectedItem == null)
-            return;
-        await _mediator.Send(new AIDetectEvent());
-    }
-
    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();

    private void ReloadThumbnailsItemClick(object sender, RoutedEventArgs e)
@@ -445,4 +480,149 @@ public partial class MainWindow
        var listItem = sender as ListViewItem;
        LvFilesContextMenu.DataContext = listItem.DataContext;
    }
+
+    private (TimeSpan Time, List<(YoloLabel Label, float Probability)> Detections)? _previousDetection;
+
+    public void AutoDetect(object sender, RoutedEventArgs e)
+    {
+        if (LvFiles.SelectedItem == null)
+            return;
+
+        _mediator.Publish(new PlaybackControlEvent(PlaybackControlEnum.Play));
+        var mediaInfo = (MediaFileInfo)LvFiles.SelectedItem;
+        _formState.CurrentMedia = mediaInfo;
+        _mediaPlayer.Stop();
+        var path = mediaInfo.Path;
+
+        var manualCancellationSource = new CancellationTokenSource();
+        var token = manualCancellationSource.Token;
+
+        _autoDetectDialog = new AutodetectDialog
+        {
+            Topmost = true,
+            Owner = this
+        };
+        _autoDetectDialog.Closing += (_, _) =>
+        {
+            manualCancellationSource.Cancel();
+            _mediaPlayer.Stop();
+        };
+        _autoDetectDialog.Top = Height - _autoDetectDialog.Height - 80;
+
+        _ = Task.Run(async () =>
+        {
+            using var detector = new YOLODetector(_config);
+            Dispatcher.Invoke(() => _autoDetectDialog.Log("Ініціалізація AI..."));
+
+            await foreach (var timeframe in _vlcFrameExtractor.ExtractFrames(path, token))
+            {
+                try
+                {
+                    var detections = _aiDetector.Detect(timeframe.Stream);
+
+                    if (!IsValidDetection(timeframe.Time, detections))
+                        continue;
+
+                    await ProcessDetection(timeframe, detections, token);
+                }
+                catch (Exception ex)
+                {
+                    _logger.LogError(ex, ex.Message);
+                    await manualCancellationSource.CancelAsync();
+                }
+            }
+            _autoDetectDialog.Close();
+        }, token);
+
+
+        _autoDetectDialog.ShowDialog();
+        Dispatcher.Invoke(() => Editor.Background = new SolidColorBrush(Color.FromArgb(1, 0, 0, 0)));
+    }
+
+    private bool IsValidDetection(TimeSpan time, List<(YoloLabel Label, float Probability)> 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(_config.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.Label.CenterX, det.Label.CenterY);
+            var closestObject = prev.Detections
+                .Select(p => new
+                {
+                    Point = p,
+                    Distance = point.SqrDistance(new Point(p.Label.CenterX, p.Label.CenterY))
+                })
+                .OrderBy(x => x.Distance)
+                .First();
+
+            // Closest object is farther than Tracking distance confidence, hence it's a different object, allow
+            if (closestObject.Distance > _config.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 + _config.AIRecognitionConfig.TrackingProbabilityIncrease)
+                return true;
+        }
+
+        return false;
+    }
+
+    private async Task ProcessDetection((TimeSpan Time, Stream Stream) timeframe, List<(YoloLabel Label, float Probability)> detections, CancellationToken token = default)
+    {
+        _previousDetection = (timeframe.Time, detections);
+        await Dispatcher.Invoke(async () =>
+        {
+            try
+            {
+                var time = timeframe.Time;
+                var labels = detections.Select(x => x.Label).ToList();
+
+                var fName = _formState.GetTimeName(timeframe.Time);
+                var imgPath = Path.Combine(_config.ImagesDirectory, $"{fName}.jpg");
+                var img = System.Drawing.Image.FromStream(timeframe.Stream);
+                img.Save(imgPath, ImageFormat.Jpeg);
+                await YoloLabel.WriteToFile(labels, Path.Combine(_config.LabelsDirectory, $"{fName}.txt"), token);
+
+                Editor.Background = new ImageBrush { ImageSource = await imgPath.OpenImage() };
+                Editor.RemoveAllAnns();
+                foreach (var (label, probability) in detections)
+                    AddAnnotationToCanvas(time, new CanvasLabel(label, Editor.RenderSize, Editor.RenderSize, probability));
+                await AddAnnotations(timeframe.Time, labels);
+
+                var log = string.Join(Environment.NewLine, detections.Select(det =>
+                    $"{_config.AnnotationClassesDict[det.Label.ClassNumber].Name}: " +
+                    $"xy=({det.Label.CenterX:F2},{det.Label.CenterY:F2}), " +
+                    $"size=({det.Label.Width:F2}, {det.Label.Height:F2}), " +
+                    $"prob: {det.Probability:F1}%"));
+                Dispatcher.Invoke(() => _autoDetectDialog.Log(log));
+
+                var thumbnailDto = await _galleryManager.CreateThumbnail(imgPath, token);
+                if (thumbnailDto != null)
+                    _datasetExplorer.AddThumbnail(thumbnailDto, labels.Select(x => x.ClassNumber));
+            }
+            catch (Exception e)
+            {
+                _logger.LogError(e, e.Message);
+            }
+        });
+    }
 }