HawkeyeVision/Plugins/LindtLeerformPlugin/ViewModels/CalibrationWindowViewModel.cs

using Avalonia.Threading;
using CommunityToolkit.Mvvm.ComponentModel;
using CommunityToolkit.Mvvm.Input;
using LindtLeerformPlugin.Services;
using OpenCvSharp;
using VisionBuilder.UI.Common.Processing;

namespace LindtLeerformPlugin.ViewModels;

public partial class CalibrationWindowViewModel : ObservableObject, IDisposable
{
    private readonly IImageSource _imageSource;
    private readonly LeerformSettings _settings;
    private CancellationTokenSource? _previewCts;
    private Mat? _lastFrame;
    private Mat? _calibCameraMatrix;
    private Mat? _calibDistCoeffs;

    [ObservableProperty] private Avalonia.Media.Imaging.Bitmap? _previewImage;
    [ObservableProperty] private string _statusText = "Ready";
    [ObservableProperty] private int _capturedImageCount;
    [ObservableProperty] private bool _isPreviewRunning;
    [ObservableProperty] private bool _isCalibrated;
    [ObservableProperty] private bool _applyCalibration;
    [ObservableProperty] private string _calibrationImageDirectory;
    [ObservableProperty] private double _rmsError;

    public CalibrationWindowViewModel(IImageSource imageSource, LeerformSettings settings)
    {
        _imageSource = imageSource;
        _settings = settings;

        _calibrationImageDirectory = settings.CalibrationImageDirectory;

        var calibrationData = CalibrationDataStore.Load();
        _isCalibrated = calibrationData is { CameraMatrix.Length: > 0 };
        _rmsError = calibrationData?.RmsError ?? 0;

        if (_isCalibrated)
            LoadCalibrationMats(calibrationData!);

        UpdateCapturedImageCount();
    }

    private void UpdateCapturedImageCount()
    {
        if (Directory.Exists(CalibrationImageDirectory))
            CapturedImageCount = Directory.GetFiles(CalibrationImageDirectory, "*.png").Length;
        else
            CapturedImageCount = 0;
    }

    [RelayCommand]
    private void StartPreview()
    {
        if (IsPreviewRunning) return;

        IsPreviewRunning = true;
        _previewCts = new CancellationTokenSource();
        _ = PreviewLoopAsync(_previewCts.Token);
        StatusText = "Preview running";
    }

    [RelayCommand]
    private void StopPreview()
    {
        if (!IsPreviewRunning) return;

        _previewCts?.Cancel();
        IsPreviewRunning = false;
        StatusText = "Preview stopped";
    }

    [RelayCommand]
    private void CaptureImage()
    {
        if (_lastFrame == null || _lastFrame.Empty()) return;

        Directory.CreateDirectory(CalibrationImageDirectory);
        var filename = $"calib_{DateTime.Now:yyyyMMdd_HHmmss_fff}.png";
        var path = Path.Combine(CalibrationImageDirectory, filename);
        Cv2.ImWrite(path, _lastFrame);
        UpdateCapturedImageCount();
        StatusText = $"Captured: {filename}";
    }

    [RelayCommand]
    private void RunCalibration()
    {
        try
        {
            StatusText = "Calibrating...";
            var patternSize = new Size(_settings.CheckerboardCols, _settings.CheckerboardRows);
            var data = CameraCalibrationService.Calibrate(CalibrationImageDirectory, patternSize);

            CalibrationDataStore.Save(data);
            _settings.CalibrationImageDirectory = CalibrationImageDirectory;

            LoadCalibrationMats(data);
            IsCalibrated = true;
            RmsError = data.RmsError;
            StatusText = $"Calibrated (RMS: {data.RmsError:F4})";
        }
        catch (Exception ex)
        {
            StatusText = $"Calibration error: {ex.Message}";
        }
    }

    [RelayCommand]
    private void ClearCalibrationImages()
    {
        if (!Directory.Exists(CalibrationImageDirectory)) return;

        foreach (var file in Directory.GetFiles(CalibrationImageDirectory, "*.png"))
            File.Delete(file);

        UpdateCapturedImageCount();
        StatusText = "Calibration images cleared";
    }

    private async Task PreviewLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            try
            {
                var frame = await _imageSource.GetImage(ct);

                _lastFrame?.Dispose();
                _lastFrame = frame;

                var displayFrame = frame;
                if (ApplyCalibration && _calibCameraMatrix != null && _calibDistCoeffs != null)
                {
                    displayFrame = new Mat();
                    Cv2.Undistort(frame, displayFrame, _calibCameraMatrix, _calibDistCoeffs);
                }

                var bitmap = ImageConverter.MatToAvaloniaBitmap(displayFrame);

                if (displayFrame != frame)
                    displayFrame.Dispose();

                await Dispatcher.UIThread.InvokeAsync(() =>
                {
                    var old = PreviewImage;
                    PreviewImage = bitmap;
                    old?.Dispose();
                });
            }
            catch (OperationCanceledException)
            {
                break;
            }
            catch (Exception ex)
            {
                await Dispatcher.UIThread.InvokeAsync(() =>
                    StatusText = $"Preview error: {ex.Message}");
                break;
            }
        }
    }

    private void LoadCalibrationMats(Models.CalibrationData data)
    {
        _calibCameraMatrix?.Dispose();
        _calibDistCoeffs?.Dispose();
        (_calibCameraMatrix, _calibDistCoeffs) = CameraCalibrationService.LoadCalibrationMats(data);
    }

    public void Dispose()
    {
        StopPreview();
        _lastFrame?.Dispose();
        _calibCameraMatrix?.Dispose();
        _calibDistCoeffs?.Dispose();
    }
}