annotations/Azaion.Common/Controls/CanvasEditor.cs

using System.Drawing;
using System.Windows;
using System.Windows.Controls;
using System.Windows.Input;
using System.Windows.Media;
using System.Windows.Media.Imaging;
using System.Windows.Shapes;
using Azaion.Common.Database;
using Azaion.Common.DTO;
using Azaion.Common.Events;
using MediatR;
using Color = System.Windows.Media.Color;
using Image = System.Windows.Controls.Image;
using Point = System.Windows.Point;
using Rectangle = System.Windows.Shapes.Rectangle;
using Size = System.Windows.Size;

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 readonly MatrixTransform _matrixTransform = new();
    private Point _panStartPoint;
    private bool _isZoomedIn;

    private const int MIN_SIZE = 12;
    private readonly TimeSpan _viewThreshold = TimeSpan.FromMilliseconds(400);

    public Image BackgroundImage { get; set; } = new() { Stretch = Stretch.Uniform };

    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)),
        };

        MouseDown += CanvasMouseDown;
        MouseMove += CanvasMouseMove;
        MouseUp += CanvasMouseUp;
        SizeChanged += CanvasResized;
        Cursor = Cursors.Cross;
        Children.Insert(0, BackgroundImage);
        Children.Add(_newAnnotationRect);
        Children.Add(_horizontalLine);
        Children.Add(_verticalLine);
        Children.Add(_classNameHint);

        Loaded += Init;
        RenderTransform = _matrixTransform;
        MouseWheel += CanvasWheel;
    }

    public void SetBackground(ImageSource? source)
    {
        SetZoom();
        BackgroundImage.Source = source;
    }

    private void SetZoom(Matrix? matrix = null)
    {
        if (matrix == null)
        {
            _matrixTransform.Matrix = Matrix.Identity;
            _isZoomedIn = false;    
        }
        else
        {
            _matrixTransform.Matrix = matrix.Value;
            _isZoomedIn = true;
        }
        // foreach (var detection in CurrentDetections) 
        //     detection.UpdateAdornerScale(scale: _matrixTransform.Matrix.M11);
    }
    
    private void CanvasWheel(object sender, MouseWheelEventArgs e)
    {
        if (Keyboard.Modifiers != ModifierKeys.Control)
            return;

        var mousePos = e.GetPosition(this);
        var scale = e.Delta > 0 ? 1.1 : 1 / 1.1;

        var matrix = _matrixTransform.Matrix;
        if (scale < 1 && matrix.M11 * scale < 1.0)
            SetZoom();
        else
        {
            matrix.ScaleAt(scale, scale, mousePos.X, mousePos.Y);
            SetZoom(matrix);
        }
    }

    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();
        if (e.LeftButton != MouseButtonState.Pressed)
            return;
        if (Keyboard.Modifiers == ModifierKeys.Control && _isZoomedIn)
        {
            _panStartPoint = e.GetPosition(this);
            SelectionState = SelectionState.PanZoomMoving;
        }
        else
            NewAnnotationStart(sender, e);
        (sender as UIElement)?.CaptureMouse();
    }

    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);

        switch (SelectionState)
        {
            case SelectionState.NewAnnCreating:
                NewAnnotationCreatingMove(sender, e);
                break;
            case SelectionState.AnnResizing:
                AnnotationResizeMove(sender, e);
                break;
            case SelectionState.AnnMoving:
                AnnotationPositionMove(sender, e);
                break;
            case SelectionState.PanZoomMoving:
                PanZoomMove(sender, e);
                break;
        }
    }

    private void PanZoomMove(object sender, MouseEventArgs e)
    {
        var currentPoint = e.GetPosition(this);
        var delta = currentPoint - _panStartPoint;

        var matrix = _matrixTransform.Matrix;
        matrix.Translate(delta.X, delta.Y);
        
        _matrixTransform.Matrix = matrix;
    }

    private void CanvasMouseUp(object sender, MouseButtonEventArgs e)
    {
        (sender as UIElement)?.ReleaseMouseCapture();
        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)
            {
                var time = GetTimeFunc();
                var control = CreateDetectionControl(CurrentAnnClass, time, new CanvasLabel
                {
                    Width = width,
                    Height = height,
                    Left = Math.Min(endPos.X, _newAnnotationStartPos.X),
                    Top = Math.Min(endPos.Y, _newAnnotationStartPos.Y),
                    Confidence = 1
                });
                control.UpdateLayout();
                CheckLabelBoundaries(control);   
            }
        }
        else if (SelectionState != SelectionState.PanZoomMoving)
            CheckLabelBoundaries(_curAnn);

        SelectionState = SelectionState.None;
        e.Handled = true;
    }

    private void CheckLabelBoundaries(DetectionControl detectionControl)
    {
        var lb = detectionControl.DetectionLabelContainer;
        var origin = lb.TranslatePoint(new Point(0, 0), this);
        lb.Children[0].Measure(new Size(double.PositiveInfinity, double.PositiveInfinity));
        var size = lb.Children[0].DesiredSize;
        var lbRect = new RectangleF((float)origin.X, (float)origin.Y, (float)size.Width, (float)size.Height);

        foreach (var c in CurrentDetections)
        {
            if (c == detectionControl)
                continue;
            var detRect = new RectangleF((float)GetLeft(c), (float)GetTop(c), (float)c.Width, (float)c.Height);
            detRect.Intersect(lbRect);


            // var intersect = detections[i].ToRectangle();
            // intersect.Intersect(detections[j].ToRectangle());

            // detectionControl.
            // var otherControls = allControls.Where(c => c != control);
            // control.UpdateLabelPosition(otherControls);
        }
    }

    private void CanvasResized(object sender, SizeChangedEventArgs e)
    {
        _horizontalLine.X2 = e.NewSize.Width;
        _verticalLine.Y2 = e.NewSize.Height;
        BackgroundImage.Width = e.NewSize.Width;
        BackgroundImage.Height = 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(Annotation annotation, List<DetectionClass> detectionClasses, Size mediaSize)
    {
        var splitTile = annotation.SplitTile;
        foreach (var detection in annotation.Detections)
        {
            var detectionClass = DetectionClass.FromYoloId(detection.ClassNumber, detectionClasses);
            CanvasLabel canvasLabel;
            if (splitTile == null)
                canvasLabel = new CanvasLabel(detection, RenderSize, mediaSize, detection.Confidence);    
            else
            {
                canvasLabel = new CanvasLabel(detection, new Size(Constants.AI_TILE_SIZE, Constants.AI_TILE_SIZE), null, detection.Confidence)
                    .ReframeFromSmall(splitTile);
                
                //From CurrentMediaSize to Render Size
                var yoloLabel = new YoloLabel(canvasLabel, mediaSize);
                canvasLabel = new CanvasLabel(yoloLabel, RenderSize, mediaSize, canvasLabel.Confidence);
            }

            CreateDetectionControl(detectionClass, annotation.Time, canvasLabel);
        }
    }

    private DetectionControl CreateDetectionControl(DetectionClass detectionClass, TimeSpan time, CanvasLabel canvasLabel)
    {
        var detectionControl = new DetectionControl(detectionClass, time, AnnotationResizeStart, canvasLabel);
        detectionControl.MouseDown += AnnotationPositionStart;
        SetLeft(detectionControl, canvasLabel.Left );
        SetTop(detectionControl, canvasLabel.Top);
        Children.Add(detectionControl);
        CurrentDetections.Add(detectionControl);
        _newAnnotationRect.Fill = new SolidColorBrush(detectionClass.Color);
        return detectionControl;
    }
    
    #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));

    public void ZoomTo(Point point)
    {
        SetZoom();
        var matrix = _matrixTransform.Matrix;
        matrix.ScaleAt(2, 2, point.X, point.Y);
        SetZoom(matrix);
    }
}