From e7afa96a0b548a24a478729859c439bfc7f7258a Mon Sep 17 00:00:00 2001
From: Alex Bezdieniezhnykh <zxsanny@gmail.com>
Date: Thu, 30 Jan 2025 12:33:24 +0200
Subject: [PATCH] fix inference UI and annotation saving

---
 Azaion.Annotator/Annotator.xaml             |   1 +
 Azaion.Annotator/Annotator.xaml.cs          | 196 ++++----------------
 Azaion.Annotator/AnnotatorEventHandler.cs   |   3 +
 Azaion.Common/Controls/DetectionControl.cs  |   2 +-
 Azaion.Common/Services/AnnotationService.cs |   2 +
 Azaion.Common/Services/InferenceService.cs  |   6 +-
 6 files changed, 47 insertions(+), 163 deletions(-)

diff --git a/Azaion.Annotator/Annotator.xaml b/Azaion.Annotator/Annotator.xaml
index 634c398..28f1160 100644
--- a/Azaion.Annotator/Annotator.xaml
+++ b/Azaion.Annotator/Annotator.xaml
@@ -475,6 +475,7 @@
                 Maximum="100" Minimum="0">
             </controls2:UpdatableProgressBar>
 
+
             <Button
                 x:Name="AIDetectBtn"
                 Grid.Column="10"
diff --git a/Azaion.Annotator/Annotator.xaml.cs b/Azaion.Annotator/Annotator.xaml.cs
index 461344e..28fbd9c 100644
--- a/Azaion.Annotator/Annotator.xaml.cs
+++ b/Azaion.Annotator/Annotator.xaml.cs
@@ -10,7 +10,6 @@ using Azaion.Annotator.DTO;
 using Azaion.Common.Database;
 using Azaion.Common.DTO;
 using Azaion.Common.DTO.Config;
-using Azaion.Common.DTO.Queue;
 using Azaion.Common.Events;
 using Azaion.Common.Extensions;
 using Azaion.Common.Services;
@@ -45,6 +44,9 @@ public partial class Annotator
     private ObservableCollection<DetectionClass> AnnotationClasses { get; set; } = new();
     private bool _suspendLayout;
 
+    public bool FollowAI = false;
+    public bool IsInferenceNow = false;
+
     private readonly TimeSpan _thresholdBefore = TimeSpan.FromMilliseconds(100);
     private readonly TimeSpan _thresholdAfter = TimeSpan.FromMilliseconds(300);
 
@@ -53,7 +55,6 @@ public partial class Annotator
     public ObservableCollection<MediaFileInfo> FilteredMediaFiles { get; set; } = new();
 
     public IntervalTree<TimeSpan, Annotation> TimedAnnotations { get; set; } = new();
-    private AutodetectDialog _autoDetectDialog = new() { Topmost = true };
 
     public Annotator(
         IConfigUpdater configUpdater,
@@ -166,7 +167,12 @@ public partial class Annotator
             }
         };
 
-        LvFiles.MouseDoubleClick += async (_, _) => await _mediator.Publish(new AnnotatorControlEvent(PlaybackControlEnum.Play));
+        LvFiles.MouseDoubleClick += async (_, _) =>
+        {
+            if (IsInferenceNow)
+                FollowAI = false;
+            await _mediator.Publish(new AnnotatorControlEvent(PlaybackControlEnum.Play));
+        };
 
         LvClasses.SelectionChanged += (_, _) =>
         {
@@ -224,6 +230,8 @@ public partial class Annotator
 
     public void OpenAnnotationResult(AnnotationResult res)
     {
+        if (IsInferenceNow)
+            FollowAI = false;
         _mediaPlayer.SetPause(true);
         Editor.RemoveAllAnns();
         _mediaPlayer.Time = (long)res.Annotation.Time.TotalMilliseconds;
@@ -399,9 +407,9 @@ public partial class Annotator
         Process.Start("explorer.exe", "/select,\"" + mediaFileInfo.Path +"\"");
     }
 
-    public void SeekTo(long timeMilliseconds)
+    public void SeekTo(long timeMilliseconds, bool setPause = true)
     {
-        _mediaPlayer.SetPause(true);
+        _mediaPlayer.SetPause(setPause);
         _mediaPlayer.Time = timeMilliseconds;
         VideoSlider.Value = _mediaPlayer.Position * 100;
 
@@ -446,6 +454,8 @@ public partial class Annotator
 
     private void PlayClick(object sender, RoutedEventArgs e)
     {
+        if (IsInferenceNow)
+            FollowAI = false;
         _mediator.Publish(new AnnotatorControlEvent(_mediaPlayer.CanPause ? PlaybackControlEnum.Pause : PlaybackControlEnum.Play));
     }
 
@@ -478,8 +488,14 @@ public partial class Annotator
 
     private (TimeSpan Time, List<Detection> Detections)? _previousDetection;
 
-    public async void AutoDetect(object sender, RoutedEventArgs e)
+    public void AutoDetect(object sender, RoutedEventArgs e)
     {
+        if (IsInferenceNow)
+        {
+            FollowAI = true;
+            return;
+        }
+
         if (LvFiles.Items.IsEmpty)
             return;
         if (LvFiles.SelectedIndex == -1)
@@ -487,28 +503,14 @@ public partial class Annotator
 
         var mct = new CancellationTokenSource();
         var token = mct.Token;
+        Dispatcher.Invoke(() => Editor.ResetBackground());
 
-        _autoDetectDialog = new AutodetectDialog
-        {
-            Topmost = true,
-            Owner = this
-        };
-        _autoDetectDialog.Closing += (_, _) =>
-        {
-            mct.Cancel();
-            _mediaPlayer.SeekTo(TimeSpan.Zero);
-            Editor.RemoveAllAnns();
-        };
-
-        _autoDetectDialog.Top = Height - _autoDetectDialog.Height - 80;
-        _autoDetectDialog.Left = 5;
-
-        _autoDetectDialog.Log("Ініціалізація AI...");
-
+        IsInferenceNow = true;
+        FollowAI = true;
         _ = Task.Run(async () =>
         {
             var mediaInfo = Dispatcher.Invoke(() => (MediaFileInfo)LvFiles.SelectedItem);
-            while (mediaInfo != null)
+            while (mediaInfo != null && !token.IsCancellationRequested)
             {
                 await Dispatcher.Invoke(async () =>
                 {
@@ -516,11 +518,11 @@ public partial class Annotator
                     await ReloadAnnotations();
                 });
 
-                await _inferenceService.RunInference(mediaInfo.Path, async (annotationImage, ct) =>
+                await _inferenceService.RunInference(mediaInfo.Path, async annotationImage =>
                 {
                     annotationImage.OriginalMediaName = mediaInfo.FName;
-                    await ProcessDetection(annotationImage, ct);
-                }, token);
+                    await ProcessDetection(annotationImage);
+                });
 
                 mediaInfo = Dispatcher.Invoke(() =>
                 {
@@ -533,156 +535,32 @@ public partial class Annotator
             }
             Dispatcher.Invoke(() =>
             {
-                _autoDetectDialog.Close();
                 _mediaPlayer.Stop();
                 LvFiles.Items.Refresh();
+                IsInferenceNow = false;
             });
         }, token);
-
-        _autoDetectDialog.ShowDialog();
-        Dispatcher.Invoke(() => Editor.ResetBackground());
     }
 
-    // private async Task DetectImage(MediaFileInfo mediaInfo, CancellationTokenSource manualCancellationSource, CancellationToken token)
-    // {
-    //     try
-    //     {
-    //         var fName = Path.GetFileNameWithoutExtension(mediaInfo.Path);
-    //         var stream = new FileStream(mediaInfo.Path, FileMode.Open);
-    //         var detections = await _aiDetector.Detect(fName, stream, token);
-    //         await ProcessDetection((TimeSpan.FromMilliseconds(0), stream), Path.GetExtension(mediaInfo.Path), 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 fName = _formState.GetTimeName(timeframe.Time);
-    //              var detections = await _aiDetector.Detect(fName, timeframe.Stream, token);
-    //
-    //              var isValid = IsValidDetection(timeframe.Time, detections);
-    //              Console.WriteLine($"Detection time: {timeframe.Time}");
-    //
-    //              var log = string.Join(Environment.NewLine, detections.Select(det =>
-    //                  $"{_appConfig.AnnotationConfig.DetectionClassesDict[det.ClassNumber].Name}: " +
-    //                  $"xy=({det.CenterX:F2},{det.CenterY:F2}), " +
-    //                  $"size=({det.Width:F2}, {det.Height:F2}), " +
-    //                  $"prob: {det.Probability:F1}%"));
-    //
-    //              log = $"Detection time: {timeframe.Time}, Valid: {isValid}. {Environment.NewLine} {log}";
-    //              Dispatcher.Invoke(() => _autoDetectDialog.Log(log));
-    //
-    //             if (timeframe.Time.TotalMilliseconds > prevSeekTime + 250)
-    //             {
-    //                 Dispatcher.Invoke(() => SeekTo(timeframe.Time));
-    //                 prevSeekTime = timeframe.Time.TotalMilliseconds;
-    //                 if (!isValid) //Show frame anyway
-    //                 {
-    //                     Dispatcher.Invoke(() =>
-    //                     {
-    //                         Editor.RemoveAllAnns();
-    //                         Editor.Background = new ImageBrush
-    //                         {
-    //                             ImageSource = timeframe.Stream.OpenImage()
-    //                         };
-    //                     });
-    //                 }
-    //             }
-    //
-    //             if (!isValid)
-    //                 continue;
-    //
-    //             mediaInfo.HasAnnotations = true;
-    //             await ProcessDetection(timeframe, ".jpg", detections, token);
-    //             await timeframe.Stream.DisposeAsync();
-    //         }
-    //         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(AnnotationImage annotationImage, CancellationToken token = default)
+    private async Task ProcessDetection(AnnotationImage annotationImage)
     {
         await Dispatcher.Invoke(async () =>
         {
             try
             {
-                var annotation = await _annotationService.SaveAnnotation(annotationImage, token);
-
-                Editor.Background = new ImageBrush { ImageSource = await annotation.ImagePath.OpenImage() };
-                Editor.RemoveAllAnns();
-                ShowAnnotations(annotation, true);
+                var annotation = await _annotationService.SaveAnnotation(annotationImage);
                 AddAnnotation(annotation);
 
+                if (FollowAI)
+                    SeekTo(annotationImage.Milliseconds, false);
+
                 var log = string.Join(Environment.NewLine, annotation.Detections.Select(det =>
                     $"{_appConfig.AnnotationConfig.DetectionClassesDict[det.ClassNumber].Name}: " +
                     $"xy=({det.CenterX:F2},{det.CenterY:F2}), " +
                     $"size=({det.Width:F2}, {det.Height:F2}), " +
-                    $"prob: {det.Probability:F1}%"));
+                    $"prob: {det.Probability*100:F0}%"));
 
-                Dispatcher.Invoke(() => _autoDetectDialog.Log(log));
+                Dispatcher.Invoke(() => StatusHelp.Text = log);
             }
             catch (Exception e)
             {
diff --git a/Azaion.Annotator/AnnotatorEventHandler.cs b/Azaion.Annotator/AnnotatorEventHandler.cs
index e0e9a56..7097a0e 100644
--- a/Azaion.Annotator/AnnotatorEventHandler.cs
+++ b/Azaion.Annotator/AnnotatorEventHandler.cs
@@ -127,8 +127,11 @@ public class AnnotatorEventHandler(
                     break;
                 case PlaybackControlEnum.Pause:
                     mediaPlayer.Pause();
+                    if (mainWindow.IsInferenceNow)
+                        mainWindow.FollowAI = false;
                     if (!mediaPlayer.IsPlaying)
                         mainWindow.BlinkHelp(HelpTexts.HelpTextsDict[HelpTextEnum.AnnotationHelp]);
+
                     if (formState.BackgroundTime.HasValue)
                     {
                         mainWindow.Editor.ResetBackground();
diff --git a/Azaion.Common/Controls/DetectionControl.cs b/Azaion.Common/Controls/DetectionControl.cs
index 0fdc09e..8128385 100644
--- a/Azaion.Common/Controls/DetectionControl.cs
+++ b/Azaion.Common/Controls/DetectionControl.cs
@@ -59,7 +59,7 @@ public class DetectionControl : Border
         };
         _probabilityLabel = new Label
         {
-            Content = probability.HasValue ? $"{probability.Value:F0}%" : string.Empty,
+            Content = probability.HasValue ? $"{probability.Value*100:F0}%" : string.Empty,
             HorizontalAlignment = HorizontalAlignment.Right,
             VerticalAlignment = VerticalAlignment.Top,
             Margin = new Thickness(0, -32, 0, 0),
diff --git a/Azaion.Common/Services/AnnotationService.cs b/Azaion.Common/Services/AnnotationService.cs
index 5e25852..387670f 100644
--- a/Azaion.Common/Services/AnnotationService.cs
+++ b/Azaion.Common/Services/AnnotationService.cs
@@ -106,6 +106,8 @@ public class AnnotationService : INotificationHandler<AnnotationsDeletedEvent>
     {
         a.Time = TimeSpan.FromMilliseconds(a.Milliseconds);
         a.Name = a.OriginalMediaName.ToTimeName(a.Time);
+        foreach (var det in a.Detections)
+            det.AnnotationName = a.Name;
         return await SaveAnnotationInner(DateTime.Now, a.OriginalMediaName, a.Time, ".jpg", a.Detections.ToList(),
             a.Source, new MemoryStream(a.Image), a.CreatedRole, a.CreatedEmail, generateThumbnail: true, cancellationToken);
     }
diff --git a/Azaion.Common/Services/InferenceService.cs b/Azaion.Common/Services/InferenceService.cs
index 0ef5efc..eed9e43 100644
--- a/Azaion.Common/Services/InferenceService.cs
+++ b/Azaion.Common/Services/InferenceService.cs
@@ -14,12 +14,12 @@ namespace Azaion.Common.Services;
 
 public interface IInferenceService
 {
-    Task RunInference(string mediaPath,  Func<AnnotationImage, CancellationToken, Task> processAnnotation, CancellationToken ct = default);
+    Task RunInference(string mediaPath,  Func<AnnotationImage, Task> processAnnotation);
 }
 
 public class PythonInferenceService(ILogger<PythonInferenceService> logger, IOptions<AIRecognitionConfig> aiConfigOptions) : IInferenceService
 {
-    public async Task RunInference(string mediaPath, Func<AnnotationImage, CancellationToken, Task> processAnnotation, CancellationToken ct = default)
+    public async Task RunInference(string mediaPath, Func<AnnotationImage, Task> processAnnotation)
     {
         using var dealer = new DealerSocket();
         var clientId = Guid.NewGuid();
@@ -42,7 +42,7 @@ public class PythonInferenceService(ILogger<PythonInferenceService> logger, IOpt
                     continue;
                 }
 
-                await processAnnotation(annotationStream, ct);
+                await processAnnotation(annotationStream);
             }
             catch (Exception e)
             {