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.Extensions;
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 };
    private RectangleF? _clampedRect;
    
    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;
        UpdateClampedRect();
    }

    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 = GetClampedPosition(e);
        _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(pos);
                break;
            case SelectionState.AnnResizing:
                AnnotationResizeMove(pos);
                break;
            case SelectionState.AnnMoving:
                AnnotationPositionMove(pos);
                e.Handled = true;
                break;
            case SelectionState.PanZoomMoving:
                PanZoomMove(pos);
                break;
        }
    }

    private Point GetClampedPosition(MouseEventArgs e)
    {
        var pos = e.GetPosition(this);
        return !_clampedRect.HasValue
            ? pos
            : new Point
            (
                Math.Clamp(pos.X, _clampedRect.Value.X, _clampedRect.Value.Right),
                Math.Clamp(pos.Y, _clampedRect.Value.Y, _clampedRect.Value.Bottom)
            );
    }

    private void PanZoomMove(Point point)
    {
        var delta = point - _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 = GetClampedPosition(e);
            _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 controlLabel = 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(controlLabel);


            // 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;
        UpdateClampedRect();
    }

    private void UpdateClampedRect()
    {
        if (BackgroundImage.Source is not BitmapSource imageSource)
        {
            _clampedRect = null;
            return;
        }

        var imgWidth = imageSource.PixelWidth;
        var imgHeight = imageSource.PixelHeight;
        var canvasWidth = ActualWidth;
        var canvasHeight = ActualHeight;

        var imgRatio = imgWidth / (double)imgHeight;
        var canvasRatio = canvasWidth / canvasHeight;

        double renderedWidth;
        double renderedHeight;

        if (imgRatio > canvasRatio)
        {
            renderedWidth = canvasWidth;
            renderedHeight = canvasWidth / imgRatio;
        }
        else
        {
            renderedHeight = canvasHeight;
            renderedWidth = canvasHeight * imgRatio;
        }
        var xOffset = (canvasWidth - renderedWidth) / 2;
        var yOffset = (canvasHeight - renderedHeight) / 2;
        
        _clampedRect = new RectangleF((float)xOffset, (float)yOffset, (float)renderedWidth, (float)renderedHeight);
    }

    #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;
        (sender as UIElement)?.CaptureMouse();
        e.Handled = true;
    }
    
    private void AnnotationResizeMove(Point point)
    {
        if (SelectionState != SelectionState.AnnResizing)
            return;
        
        var x = GetLeft(_curAnn);
        var y = GetTop(_curAnn);
        var offsetX = point.X - _lastPos.X;
        var offsetY = point.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 = point;
    }

    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(Point point)
    {
        if (SelectionState != SelectionState.AnnMoving)
            return;
        
        var offsetX = point.X - _lastPos.X;
        var offsetY = point.Y - _lastPos.Y;

        var nextLeft = GetLeft(_curAnn) + offsetX;
        var nextTop = GetTop(_curAnn) + offsetY;

        if (_clampedRect.HasValue)
        {
            nextLeft = Math.Clamp(nextLeft, _clampedRect.Value.X, _clampedRect.Value.Right - _curAnn.Width);
            nextTop = Math.Clamp(nextTop, _clampedRect.Value.Y, _clampedRect.Value.Bottom - _curAnn.Height);
        }

        SetLeft(_curAnn, nextLeft);
        SetTop(_curAnn, nextTop);
        _lastPos = point;
    }
    
    #endregion
    
    #region NewAnnotation

    private void NewAnnotationStart(object sender, MouseButtonEventArgs e)
    {
        _newAnnotationStartPos = e.GetPosition(this);
        SetLeft(_newAnnotationRect, _newAnnotationStartPos.X);
        SetTop(_newAnnotationRect, _newAnnotationStartPos.Y);
        _newAnnotationRect.MouseMove += (sender, e) =>
        {
            var currentPos = e.GetPosition(this);
            NewAnnotationCreatingMove(currentPos);
        };

        SelectionState = SelectionState.NewAnnCreating;
    }
    
    private void NewAnnotationCreatingMove(Point point)
    {
        if (SelectionState != SelectionState.NewAnnCreating)
            return;
        
        var diff = point - _newAnnotationStartPos;
        
        _newAnnotationRect.Height = Math.Abs(diff.Y);
        _newAnnotationRect.Width = Math.Abs(diff.X);
        
        if (diff.X < 0)
            SetLeft(_newAnnotationRect, point.X);
        if (diff.Y < 0)
            SetTop(_newAnnotationRect, point.Y);
    }

    public void CreateDetections(Annotation annotation, List<DetectionClass> detectionClasses, Size mediaSize)
    {
        foreach (var detection in annotation.Detections)
        {
            var detectionClass = DetectionClass.FromYoloId(detection.ClassNumber, detectionClasses);
            CanvasLabel canvasLabel;
            if (!annotation.IsSplit || mediaSize.FitSizeForAI())
                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(annotation.SplitTile!);
                
                //From CurrentMediaSize to Render Size
                var yoloLabel = new YoloLabel(canvasLabel, mediaSize);
                canvasLabel = new CanvasLabel(yoloLabel, RenderSize, mediaSize, canvasLabel.Confidence);
            }

            var control = CreateDetectionControl(detectionClass, annotation.Time, canvasLabel);
            control.UpdateLayout();
            CheckLabelBoundaries(control);
        }
    }

    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 ZoomTo(Point point)
    {
        SetZoom();
        var matrix = _matrixTransform.Matrix;
        matrix.ScaleAt(2, 2, point.X, point.Y);
        SetZoom(matrix);
    }
}