annotations/Azaion.Common/Controls/CanvasEditor.cs

using System.Windows;
using System.Windows.Controls;
using System.Windows.Input;
using System.Windows.Media;
using System.Windows.Shapes;
using Azaion.Annotator.DTO;
using Azaion.Common.DTO;
using MediatR;
using Color = System.Windows.Media.Color;
using Rectangle = System.Windows.Shapes.Rectangle;

namespace Azaion.Common.Controls;

public class CanvasEditor : Canvas
{
    private Point _lastPos;
    public SelectionState SelectionState { get; set; } = SelectionState.None;

    private readonly Rectangle _newAnnotationRect;
    private Point _newAnnotationStartPos;
    
    private readonly Line _horizontalLine;
    private readonly Line _verticalLine;
    private readonly TextBlock _classNameHint;
    
    private Rectangle _curRec = new();
    private DetectionControl _curAnn = null!;
    
    private const int MIN_SIZE = 20;
    private readonly TimeSpan _viewThreshold = TimeSpan.FromMilliseconds(400);
    
    public IMediator Mediator { get; set; } = null!;

    public static readonly DependencyProperty GetTimeFuncProp =
        DependencyProperty.Register(
            nameof(GetTimeFunc),
            typeof(Func<TimeSpan>),
            typeof(CanvasEditor),
            new PropertyMetadata(null));

    public Func<TimeSpan> GetTimeFunc
    {
        get => (Func<TimeSpan>)GetValue(GetTimeFuncProp);
        set => SetValue(GetTimeFuncProp, value);
    }

    private DetectionClass _currentAnnClass = null!;
    public DetectionClass CurrentAnnClass
    {
        get => _currentAnnClass; 
        set
        {
            _verticalLine.Stroke = value.ColorBrush;
            _verticalLine.Fill = value.ColorBrush;
            _horizontalLine.Stroke = value.ColorBrush;
            _horizontalLine.Fill = value.ColorBrush;
            _classNameHint.Text = value.ShortName;
            
            _newAnnotationRect.Stroke = value.ColorBrush;
            _newAnnotationRect.Fill = value.ColorBrush;
            _currentAnnClass = value;
        }
    }

    public readonly List<DetectionControl> CurrentDetections = new();
    
    public CanvasEditor()
    {
        _horizontalLine = new Line
        {
            HorizontalAlignment = HorizontalAlignment.Stretch,
            Stroke = new SolidColorBrush(Colors.Blue),
            Fill = new SolidColorBrush(Colors.Blue),
            StrokeDashArray = [5],
            StrokeThickness = 2
        };
        _verticalLine = new Line
        {
            VerticalAlignment = VerticalAlignment.Stretch,
            Stroke = new SolidColorBrush(Colors.Blue),
            Fill = new SolidColorBrush(Colors.Blue),
            StrokeDashArray = [5],
            StrokeThickness = 2
        };
        _classNameHint = new TextBlock
        {
            Text = CurrentAnnClass?.ShortName ?? "",
            Foreground = new SolidColorBrush(Colors.Black),
            Cursor = Cursors.Arrow,
            FontSize = 16,
            FontWeight = FontWeights.Bold
        };
        _newAnnotationRect = new Rectangle
        {
            Name = "selector",
            Height = 0,
            Width = 0,
            Stroke = new SolidColorBrush(Colors.Gray),
            Fill = new SolidColorBrush(Color.FromArgb(128, 128, 128, 128)),
        };

        KeyDown += (_, args) =>
        {
            Console.WriteLine($"pressed {args.Key}");
        };
        MouseDown += CanvasMouseDown;
        MouseMove += CanvasMouseMove;
        MouseUp += CanvasMouseUp;
        SizeChanged += CanvasResized;
        Cursor = Cursors.Cross;

        Children.Add(_newAnnotationRect);
        Children.Add(_horizontalLine);
        Children.Add(_verticalLine);
        Children.Add(_classNameHint);

        Loaded += Init;
    }

    private void Init(object sender, RoutedEventArgs e)
    {
        _horizontalLine.X1 = 0;
        _horizontalLine.X2 = ActualWidth;
        _verticalLine.Y1 = 0;
        _verticalLine.Y2 = ActualHeight;
    }

    private void CanvasMouseDown(object sender, MouseButtonEventArgs e)
    {
        ClearSelections();
        NewAnnotationStart(sender, e);
    }

    private void CanvasMouseMove(object sender, MouseEventArgs e)
    {
        var pos = e.GetPosition(this);
        _horizontalLine.Y1 = _horizontalLine.Y2 = pos.Y;
        _verticalLine.X1 = _verticalLine.X2 = pos.X;
        SetLeft(_classNameHint, pos.X + 10);
        SetTop(_classNameHint, pos.Y - 30);
        
        if (e.LeftButton != MouseButtonState.Pressed)
            return;
        if (SelectionState == SelectionState.NewAnnCreating)
            NewAnnotationCreatingMove(sender, e);
        
        if (SelectionState == SelectionState.AnnResizing)
            AnnotationResizeMove(sender, e);

        if (SelectionState == SelectionState.AnnMoving)
            AnnotationPositionMove(sender, e);
    }

    private void CanvasMouseUp(object sender, MouseButtonEventArgs e)
    {
        if (SelectionState == SelectionState.NewAnnCreating)
        {
            var endPos = e.GetPosition(this);
            _newAnnotationRect.Width = 0;
            _newAnnotationRect.Height = 0;
            var width = Math.Abs(endPos.X - _newAnnotationStartPos.X);
            var height = Math.Abs(endPos.Y - _newAnnotationStartPos.Y);
            if (width < MIN_SIZE || height < MIN_SIZE)
                return;

            var time = GetTimeFunc();
            CreateDetectionControl(CurrentAnnClass, time, new CanvasLabel
            {
                Width = width,
                Height = height,
                X = Math.Min(endPos.X, _newAnnotationStartPos.X),
                Y = Math.Min(endPos.Y, _newAnnotationStartPos.Y),
                Confidence = 1
            });
        }

        SelectionState = SelectionState.None;
        e.Handled = true;
    }
    
    private void CanvasResized(object sender, SizeChangedEventArgs e)
    {
        _horizontalLine.X2 = e.NewSize.Width;
        _verticalLine.Y2 = e.NewSize.Height;
    }
    
    #region Annotation Resizing & Moving
    
    private void AnnotationResizeStart(object sender, MouseEventArgs e)
    {
        SelectionState = SelectionState.AnnResizing;
        _lastPos = e.GetPosition(this);
        _curRec = (Rectangle)sender;
        _curAnn = (DetectionControl)((Grid)_curRec.Parent).Parent;
        e.Handled = true;
    }
    
    private void AnnotationResizeMove(object sender, MouseEventArgs e)
    {
        if (SelectionState != SelectionState.AnnResizing)
            return;
        
        var currentPos = e.GetPosition(this);

        var x = GetLeft(_curAnn);
        var y = GetTop(_curAnn);
        var offsetX = currentPos.X - _lastPos.X;
        var offsetY = currentPos.Y - _lastPos.Y;
        switch (_curRec.HorizontalAlignment, _curRec.VerticalAlignment)
        {
            case (HorizontalAlignment.Left, VerticalAlignment.Top):
                _curAnn.Width = Math.Max(MIN_SIZE, _curAnn.Width - offsetX);
                _curAnn.Height = Math.Max(MIN_SIZE, _curAnn.Height - offsetY);
                SetLeft(_curAnn, x + offsetX);
                SetTop(_curAnn, y + offsetY);
                break;
            case (HorizontalAlignment.Center, VerticalAlignment.Top):
                _curAnn.Height = Math.Max(MIN_SIZE, _curAnn.Height - offsetY);
                SetTop(_curAnn, y + offsetY);
                break;
            case (HorizontalAlignment.Right, VerticalAlignment.Top):
                _curAnn.Width = Math.Max(MIN_SIZE, _curAnn.Width + offsetX);
                _curAnn.Height = Math.Max(MIN_SIZE, _curAnn.Height - offsetY);
                SetTop(_curAnn, y + offsetY);
                break;
            
            case (HorizontalAlignment.Left, VerticalAlignment.Center):
                _curAnn.Width = Math.Max(MIN_SIZE, _curAnn.Width - offsetX);
                SetLeft(_curAnn, x + offsetX);
                break;
            case (HorizontalAlignment.Right, VerticalAlignment.Center):
                _curAnn.Width = Math.Max(MIN_SIZE, _curAnn.Width + offsetX);
                break;
            
            case (HorizontalAlignment.Left, VerticalAlignment.Bottom):
                _curAnn.Width = Math.Max(MIN_SIZE, _curAnn.Width - offsetX);
                _curAnn.Height = Math.Max(MIN_SIZE, _curAnn.Height + offsetY);
                SetLeft(_curAnn, x + offsetX);
                break;
            case (HorizontalAlignment.Center, VerticalAlignment.Bottom):
                _curAnn.Height = Math.Max(MIN_SIZE, _curAnn.Height + offsetY);
                break;
            case (HorizontalAlignment.Right, VerticalAlignment.Bottom):
                _curAnn.Width = Math.Max(MIN_SIZE, _curAnn.Width + offsetX);
                _curAnn.Height = Math.Max(MIN_SIZE, _curAnn.Height + offsetY);
                break;
        }
        _lastPos = currentPos;
    }

    private void AnnotationPositionStart(object sender, MouseEventArgs e)
    {
        _lastPos = e.GetPosition(this);
        _curAnn = (DetectionControl)sender;

        if (!Keyboard.IsKeyDown(Key.LeftCtrl) && !Keyboard.IsKeyDown(Key.RightCtrl)) 
            ClearSelections();

        _curAnn.IsSelected = true;
        
        SelectionState = SelectionState.AnnMoving;
        e.Handled = true;
    }
    
    private void AnnotationPositionMove(object sender, MouseEventArgs e)
    {
        if (SelectionState != SelectionState.AnnMoving)
            return;
        
        var currentPos = e.GetPosition(this);
        var offsetX = currentPos.X - _lastPos.X;
        var offsetY = currentPos.Y - _lastPos.Y;
        
        SetLeft(_curAnn, GetLeft(_curAnn) + offsetX);
        SetTop(_curAnn, GetTop(_curAnn) + offsetY);
        _lastPos = currentPos;
        e.Handled = true;
    }
    
    #endregion
    
    #region NewAnnotation

    private void NewAnnotationStart(object sender, MouseButtonEventArgs e)
    {
        _newAnnotationStartPos = e.GetPosition(this);
        
        SetLeft(_newAnnotationRect, _newAnnotationStartPos.X);
        SetTop(_newAnnotationRect, _newAnnotationStartPos.Y);
        _newAnnotationRect.MouseMove += NewAnnotationCreatingMove;
        
        SelectionState = SelectionState.NewAnnCreating;
    }
    
    private void NewAnnotationCreatingMove(object sender, MouseEventArgs e)
    {
        if (SelectionState != SelectionState.NewAnnCreating)
            return;
        
        var currentPos = e.GetPosition(this);
        var diff = currentPos - _newAnnotationStartPos;
        
        _newAnnotationRect.Height = Math.Abs(diff.Y);
        _newAnnotationRect.Width = Math.Abs(diff.X);
        
        if (diff.X < 0)
            SetLeft(_newAnnotationRect, currentPos.X);
        if (diff.Y < 0)
            SetTop(_newAnnotationRect, currentPos.Y);
    }

    public void CreateDetections(TimeSpan time, IEnumerable<Detection> detections, List<DetectionClass> detectionClasses, Size videoSize)
    {
        foreach (var detection in detections)
        {
            var detectionClass = DetectionClass.FromYoloId(detection.ClassNumber, detectionClasses);
            var canvasLabel = new CanvasLabel(detection, RenderSize, videoSize, detection.Confidence);
            CreateDetectionControl(detectionClass, time, canvasLabel);
        }
    }

    private void CreateDetectionControl(DetectionClass detectionClass, TimeSpan time, CanvasLabel canvasLabel)
    {
        var detectionControl = new DetectionControl(detectionClass, time, AnnotationResizeStart, canvasLabel);
        detectionControl.MouseDown += AnnotationPositionStart;
        SetLeft(detectionControl, canvasLabel.X );
        SetTop(detectionControl, canvasLabel.Y);
        Children.Add(detectionControl);
        CurrentDetections.Add(detectionControl);
        _newAnnotationRect.Fill = new SolidColorBrush(detectionClass.Color);
    }
    
    #endregion

    private void RemoveAnnotations(IEnumerable<DetectionControl> listToRemove)
    {
        foreach (var ann in listToRemove)
        {
            Children.Remove(ann);
            CurrentDetections.Remove(ann);
        }   
    }

    public void RemoveAllAnns()
    {
        foreach (var ann in CurrentDetections)
            Children.Remove(ann);
        CurrentDetections.Clear();
    }

    public void RemoveSelectedAnns() => RemoveAnnotations(CurrentDetections.Where(x => x.IsSelected).ToList());

    private void ClearSelections()
    {
        foreach (var ann in CurrentDetections)
            ann.IsSelected = false;
    }

    public void ClearExpiredAnnotations(TimeSpan time)
    {
        var expiredAnns = CurrentDetections.Where(x =>
            Math.Abs((time - x.Time).TotalMilliseconds) > _viewThreshold.TotalMilliseconds)
            .ToList();
        RemoveAnnotations(expiredAnns);
    }

    public void ResetBackground() => Background = new SolidColorBrush(Color.FromArgb(1, 0, 0, 0));
}