Files
annotations/Azaion.Annotator/MainWindow.xaml.cs
T
Alex Bezdieniezhnykh 418a2116b7 add autodetection
2024-11-02 13:09:00 +02:00

629 lines
25 KiB
C#

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 Size = System.Windows.Size;
using IntervalTree;
using Microsoft.Extensions.Logging;
using OpenTK.Graphics.OpenGL;
using Serilog;
using MediaPlayer = LibVLCSharp.Shared.MediaPlayer;
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 readonly VLCFrameExtractor _vlcFrameExtractor;
private readonly IAIDetector _aiDetector;
private 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 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();
private AutodetectDialog _autoDetectDialog;
public MainWindow(LibVLC libVLC, MediaPlayer mediaPlayer,
IMediator mediator,
FormState formState,
IConfigRepository configRepository,
HelpWindow helpWindow,
DatasetExplorer datasetExplorer,
ILogger<MainWindow> logger,
IGalleryManager galleryManager,
VLCFrameExtractor vlcFrameExtractor,
IAIDetector aiDetector)
{
InitializeComponent();
_libVLC = libVLC;
_mediaPlayer = mediaPlayer;
_mediator = mediator;
_formState = formState;
_configRepository = configRepository;
_config = _configRepository.Get();
_helpWindow = helpWindow;
_datasetExplorer = datasetExplorer;
_logger = logger;
_galleryManager = galleryManager;
_vlcFrameExtractor = vlcFrameExtractor;
_aiDetector = aiDetector;
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 () => await _galleryManager.RefreshThumbnails());
_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();
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;
}
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)
AddAnnotationToCanvas(time, new CanvasLabel(ann, Editor.RenderSize, _formState.CurrentVideoSize));
}
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();
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);
await AddAnnotations(time, await YoloLabel.ReadFromFile(file.FullName, cancellationToken));
}
}
public async Task AddAnnotations(TimeSpan? time, List<YoloLabel> annotations)
{
var timeValue = time ?? TimeSpan.FromMinutes(0);
var previousAnnotations = Annotations.Query(timeValue);
Annotations.Remove(previousAnnotations);
Annotations.Add(timeValue.Subtract(_thresholdBefore), timeValue.Add(_thresholdAfter), annotations);
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, new AnnotationResult(timeValue, _formState.GetTimeName(time), annotations, _config));
await File.WriteAllTextAsync($"{_config.ResultsDirectory}/{_formState.VideoName}.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 =>
{
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(_config.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(_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]);
DataContext = this;
}
private void OnFormClosed(object? sender, EventArgs e)
{
_mediaPlayer.Stop();
_mediaPlayer.Dispose();
_libVLC.Dispose();
_configRepository.Save(_config);
Application.Current.Shutdown();
}
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))
{
_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 Thumb_OnDragCompleted(object sender, DragCompletedEventArgs e) => _ = SaveUserSettings();
private void ReloadThumbnailsItemClick(object sender, RoutedEventArgs e)
{
var result = MessageBox.Show($"Видалити всі іконки і згенерувати нову базу іконок в {_config.ThumbnailsDirectory}?",
"Підтвердження оновлення іконок", MessageBoxButton.YesNo, MessageBoxImage.Question);
if (result != MessageBoxResult.Yes)
return;
_galleryManager.ClearThumbnails();
_galleryManager.RefreshThumbnails();
}
private void LvFilesContextOpening(object sender, ContextMenuEventArgs e)
{
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);
}
});
}
}