GSS2Rework/GSS2/Commands/CameraTuning.cs

using System.CommandLine;

using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;

using OpenCvSharp;

using GSS2.Core;
using GSS2.Core.Logging;
using GSS2.Core.Hardware;
using GSS2.Core.Extensions;

namespace GSS2.Commands;

public static class CameraTuning
{
    private class CameraTuningWorker : BackgroundService
    {
        private const double PatternOuterRadius = 180.0;
        private const double PatternCenterRadius = 60.0;

        private static readonly Scalar[] KnownSegments =
        {
            new Scalar(0,255,255),   // Голубой
            new Scalar(0,0,255),     // Синий
            new Scalar(255,0,255),   // Пурпурный
            new Scalar(128,128,128), // Серый
            new Scalar(255,0,0),     // Красный
            new Scalar(255,128,0),   // Оранжевый
            new Scalar(255,255,0),   // Жёлтый
            new Scalar(0,255,0)      // Зелёный
        };

        private readonly ILogger<CameraTuningWorker> _logger;
        private readonly IHostApplicationLifetime _applicationLifetime;
        private readonly CameraService _cameraService;
        private readonly IlluminatorService _illuminatorService;
        private readonly FileInfo? _output;
        private double _gainR;
        private double _gainG;
        private double _gainB;
        private readonly double _threshold;
        private readonly uint _steps;
        private readonly double _intensityWhite;

        public CameraTuningWorker(
            ILogger<CameraTuningWorker> logger,
            IHostApplicationLifetime applicationLifetime,
            CameraService cameraService,
            IlluminatorService illuminatorService,
            FileInfo? output,
            double gainR,
            double gainG,
            double gainB,
            double threshold,
            uint steps,
            double intensityWhite
        )
        {
            _applicationLifetime = applicationLifetime;
            _logger = logger;
            _cameraService = cameraService;
            _illuminatorService = illuminatorService;
            _output = output;
            _gainR = gainR;
            _gainG = gainG;
            _gainB = gainB;
            _threshold = threshold;
            _steps = steps;
            _intensityWhite = intensityWhite;
        }

        protected override async Task ExecuteAsync(CancellationToken stoppingToken)
        {
            // Логгируем информацию
            _logger.LogInformation("Регулировка параметров камеры...");
            _logger.LogInformation("Файл выходного изображения: {}", _output);
            _logger.LogInformation("Начальное усиление цвета (красный канал): {}", _gainR);
            _logger.LogInformation("Начальное усиление цвета (зелёный канал): {}", _gainG);
            _logger.LogInformation("Начальное усиление цвета (синий канал): {}", _gainB);
            _logger.LogInformation("Порог прекращения регулировки: {}", _threshold);
            _logger.LogInformation("Максимальное количество шагов регулировки: {}", _steps);
            _logger.LogInformation("Интенсивность белого канала осветителя: {}", _intensityWhite);

            // Инициализируем камеру
            await _cameraService.InitializeAsync(stoppingToken);
            stoppingToken.ThrowIfCancellationRequested();

            Mat? image;
            bool optimalGainsFound = false;

            // Включаем осветитель
            _illuminatorService.SetIntensity(_intensityWhite, 0, 0);
            _illuminatorService.TurnOn();

            // Пока количество шагов не дойдёт до максимума или не будут найдены оптимальные параметры
            for (int step = 0; step < _steps || !optimalGainsFound; step++)
            {
                _logger.LogInformation("Шаг {}/{}. Параметры баланса белого:\n\tКрасный: {}\n\tЗелёный: {}\n\tСиний: {}", step + 1, _steps, _gainR, _gainG, _gainB);

                // Устанавливаем усиления цветовых каналов в конфигурацию камеры
                _cameraService.Configuration.DecodeGainR = _gainR;
                _cameraService.Configuration.DecodeGainG = _gainG;
                _cameraService.Configuration.DecodeGainB = _gainB;

                // Захватываем изображение
                try
                {
                    image = await _cameraService.CaptureImage(stoppingToken, 2);

                    if (image is null)
                        throw new Exception("Не удаётся захватить изображение с камеры.");

                    // Проверяем тип изображения
                    if (image.Type() != MatType.CV_8UC3 &&
                        image.Type() != MatType.CV_8UC4 &&
                        image.Type() != MatType.CV_16UC3 &&
                        image.Type() != MatType.CV_16UC4)
                        _logger.LogCritical("Захваченное изображение имеет неподдерживаемый формат: {}", image.Type().ToString());

                    // Приводим изображение к типу CV_8
                    if (image.Type() == MatType.CV_16UC3)
                        image.ConvertTo(image, MatType.CV_8UC3, 1 / 256.0);
                    if (image.Type() == MatType.CV_16UC4)
                        image.ConvertTo(image, MatType.CV_8UC4, 1 / 256.0);

                    // Если есть альфа-канал, то удаляем его
                    if (image.Channels() == 4)
                        Cv2.CvtColor(image, image, ColorConversionCodes.BGRA2BGR);
                }
                catch (Exception ex)
                {
                    _logger.LogCritical(ex, "Ошибка при захвате изображения");
                    throw;
                }

                // Если нужно, сохраняем изображение в файл
                try
                {
                    if (_output is not null)
                    {
                        image.ImWrite(path: _output.FullName);
                        _logger.LogInformation("Изображение сохранено: {}", _output.FullName);
                    }
                }
                catch (Exception ex)
                {
                    // Сохранение в файл - в данном случае не критическая операция,
                    // поэтому логгируем сбой как некритическую ошибку и не прерываем программу
                    _logger.LogError(ex, "Ошибка при сохранении изображения");
                }

                try
                {
                    var gains = ComputeColorGains(image);

                    if (gains is null)
                        throw new Exception("Не удаётся вычислить параметры баланса белого");

                    var (r, g, b) = gains.Value;

                    // Если усиления по всем каналам изменились меньше чем на пороговое значение,
                    // то мы считаем, что оптимальные параметры баланса белого найдены, и прерываем цикл
                    if (Math.Abs(_gainR - r) < _threshold &&
                        Math.Abs(_gainG - g) < _threshold &&
                        Math.Abs(_gainB - b) < _threshold)
                    {
                        optimalGainsFound = true;
                        break;
                    }

                    // Если оптимальные параметры баланса белого не найдены,
                    // то определяем новое усиление, как среднее значение
                    // между усилением с прошлого шага
                    // и усилением, найденным на этом шаге
                    _gainR = (_gainR + r) / 2;
                    _gainG = (_gainG + g) / 2;
                    _gainB = (_gainB + b) / 2;
                }
                catch (Exception ex)
                {
                    _logger.LogCritical(ex, "Exception occurred while gains computing");
                    break;
                }
                step++;
            }

            _illuminatorService.TurnOff();

            if (!optimalGainsFound)
                _logger.LogWarning("Оптимальные параметры баланса белого не найдены!");
            else
                _logger.LogInformation("Оптимальные параметры баланса белого:\n\tКрасный: {}\n\tЗелёный: {}\n\tСиний: {}", _gainR, _gainG, _gainB);

            _applicationLifetime.StopApplication();
        }

        // TODO: Добавить комментарии к методам вычисления параметров баланса белого
        /// <summary>
        /// Вычисляет параметры баланса белого
        /// </summary>
        /// <param name="image">Изображение opencv в формате CV_8UC3</param>
        /// <returns>
        /// <see langword="null"/> если по какой-то причине вычислить параметры баланса белого не удалось,
        /// иначе вычисленные параметры баланса белого
        /// </returns>
        private static (double r, double g, double b)? ComputeColorGains(Mat image)
        {
            var circle = FindCircle(image);
            if (circle is null)
                return null;

            var (rx, ry, rr) = circle.Value;

            double scale = rr / PatternOuterRadius;
            double innerPx = PatternCenterRadius * scale;
            int midR = (int)((innerPx + rr) / 2.0);

            double phi0 = DetectOrientation(image, rx, ry, rr);
            if (double.IsNaN(phi0))
                phi0 = Math.PI / 2.0;

            var samples = new List<Vec3d>();
            var refs = new List<Vec3d>();

            int N = KnownSegments.Length;

            for (int i = 0; i < N; i++)
            {
                if (i == 3)
                    continue;

                double theta = phi0 - (2 * Math.PI / N) * i;

                int cx = (int)(rx + midR * Math.Cos(theta));
                int cy = (int)(ry - midR * Math.Sin(theta));

                var roi = new Rect(
                    Math.Max(cx - 8, 0),
                    Math.Max(cy - 8, 0),
                    Math.Min(16, image.Width - Math.Max(cx - 8, 0)),
                    Math.Min(16, image.Height - Math.Max(cy - 8, 0))
                );

                if (roi.Width < 8 || roi.Height < 8)
                    continue;

                using var patch = new Mat(image, roi);
                var mean = Cv2.Mean(patch);

                // Поскольку изображение имеет формат BGR, то изменяем порядок каналов под RGB
                samples.Add(new Vec3d(mean.Val2, mean.Val1, mean.Val0));

                var refColor = KnownSegments[i];
                refs.Add(new Vec3d(refColor.Val2, refColor.Val1, refColor.Val0));
            }

            if (samples.Count < 3)
                return null;

            // Формируем матрицы для least squares
            var S = new Mat(samples.Count, 3, MatType.CV_64F);
            var R = new Mat(samples.Count, 3, MatType.CV_64F);

            for (int i = 0; i < samples.Count; i++)
            {
                S.Set(i, 0, samples[i][0]);
                S.Set(i, 1, samples[i][1]);
                S.Set(i, 2, samples[i][2]);

                R.Set(i, 0, refs[i][0]);
                R.Set(i, 1, refs[i][1]);
                R.Set(i, 2, refs[i][2]);
            }

            var M = new Mat();
            Cv2.Solve(S, R, M, DecompTypes.Normal);

            double r = Math.Max(M.At<double>(0, 0), 1.0);
            double g = 1.0;
            double b = Math.Max(M.At<double>(2, 2), 1.0);

            return (r, g, b);
        }
        private static (int x, int y, int r)? FindCircle(Mat image)
        {
            using var gray = new Mat();
            Cv2.CvtColor(image, gray, ColorConversionCodes.BGR2GRAY);
            Cv2.GaussianBlur(gray, gray, new Size(9, 9), 2);

            var circles = Cv2.HoughCircles(
                gray,
                HoughModes.Gradient,
                dp: 1.2,
                minDist: gray.Rows / 3,
                param1: 100,
                param2: 30,
                minRadius: (int)(gray.Rows * 0.3),
                maxRadius: (int)(gray.Rows * 0.6));

            if (circles.Length == 0)
                return null;

            var c = circles[0];
            return ((int)c.Center.X, (int)c.Center.Y, (int)c.Radius);
        }
        private static double DetectOrientation(Mat image, int rx, int ry, int rr)
        {
            double scale = rr / PatternOuterRadius;
            int innerR = (int)(PatternCenterRadius * scale);

            using var gray = new Mat();
            Cv2.CvtColor(image, gray, ColorConversionCodes.BGR2GRAY);

            using var mask = Mat.Zeros(image.Size(), MatType.CV_8U).ToMat();
            Cv2.Circle(mask, new Point(rx, ry), innerR, Scalar.White, -1);

            var whitePoints = new List<Point>();

            for (int y = 0; y < gray.Rows; y++)
                for (int x = 0; x < gray.Cols; x++)
                    if (mask.At<byte>(y, x) == 255 && gray.At<byte>(y, x) > 128)
                        whitePoints.Add(new Point(x, y));

            if (whitePoints.Count < 20)
                return double.NaN;

            double cx = whitePoints.Average(p => p.X);
            double cy = whitePoints.Average(p => p.Y);

            return Math.Atan2(ry - cy, cx - rx);
        }
    }

    public static readonly RootCommand RootCommand;
    private static readonly Command Command;
    private static readonly Option<FileInfo> Output;
    private static readonly Option<double> GainR;
    private static readonly Option<double> GainG;
    private static readonly Option<double> GainB;
    private static readonly Option<double> Threshold;
    private static readonly Option<uint> Steps;
    private static readonly Option<double> IntensityWhite;

    static CameraTuning()
    {
        // Создаём команды, опции, аргументы и пр.
        RootCommand = new("Запустить регулировку параметров камеры");
        Command = new("camera-tuning")
        {
            Description = "Запустить регулировку параметров камеры"
        };
        Output = new("--output", ["-o"])
        {
            Description = "Файл выходного изображения",
            Required = false
        };
        GainR = new("--gain-r", "-r")
        {
            Description = "Начальное усиление цвета (красный канал)",
            DefaultValueFactory = (_) => 1
        };
        GainG = new("--gain-g", "-g")
        {
            Description = "Начальное усиление цвета (зелёный канал)",
            DefaultValueFactory = (_) => 1
        };
        GainB = new("--gain-b", "-b")
        {
            Description = "Начальное усиление цвета (синий канал)",
            DefaultValueFactory = (_) => 1
        };
        Threshold = new("--threshold", "-t")
        {
            Description = "Порог прекращения регулировки",
            DefaultValueFactory = (_) => 0.1
        };
        Steps = new("--steps", "-s")
        {
            Description = "Максимальное количество шагов регулировки",
            DefaultValueFactory = (_) => 10
        };
        IntensityWhite = new("--intensity-white", ["-iw"])
        {
            Description = "Интенсивность белого канала осветителя (от 0 до 1)",
            DefaultValueFactory = (_) => 1
        };

        // Изменяем описания стандартных опций
        Helper.TranslateDefaultOptionDescriptions(RootCommand);

        // Наполняем команды
        RootCommand.Add(Output);
        RootCommand.Add(GainR);
        RootCommand.Add(GainG);
        RootCommand.Add(GainB);
        RootCommand.Add(Threshold);
        RootCommand.Add(Steps);
        RootCommand.Add(IntensityWhite);
        RootCommand.SetAction(CommandAction);

        Command.Add(Output);
        Command.Add(GainR);
        Command.Add(GainG);
        Command.Add(GainB);
        Command.Add(Threshold);
        Command.Add(Steps);
        Command.Add(IntensityWhite);
        Command.SetAction(CommandAction);
    }

    /// <summary>
    /// Получить команду как корневую
    /// </summary>
    public static RootCommand GetCommand() => RootCommand;

    /// <summary>
    /// Зарегистрировать команду как подкоманду
    /// </summary>
    /// <param name="command">Команда в которой необходимо зарегистрировать подкоманду</param>
    public static void RegisterCommand(Command command) => command.Add(Command);

    private static void CommandAction(ParseResult result)
    {
        /// Получаем значения аргументов командной строки
        var output = result.GetValue(Output);
        var gainR = result.GetRequiredValue(GainR);
        var gainG = result.GetRequiredValue(GainG);
        var gainB = result.GetRequiredValue(GainB);
        var threshold = result.GetRequiredValue(Threshold);
        var steps = result.GetRequiredValue(Steps);
        var intensityWhite = result.GetRequiredValue(IntensityWhite);

        // Создаём хост
        var builder = Host.CreateApplicationBuilder();

        // Добавляем конфигурацию из appsettings.json
        builder.Configuration.AddAppSettingsJson();

        // Конфигурируем логгер
        builder.Logging.ConfigureLogging();
        builder.Services.AddSingleton<LibCameraLogSink>();

        // Добавляем сервисы аппаратной части
        builder.Services.AddIlluminatorService("Hardware:Illuminator");
        builder.Services.AddCameraService("Hardware:Camera");

        // Добавляем рабочий сервис
        builder.Services.AddHostedService<CameraTuningWorker>(services =>
            new CameraTuningWorker(
                services.GetRequiredService<ILogger<CameraTuningWorker>>(),
                services.GetRequiredService<IHostApplicationLifetime>(),
                services.GetRequiredService<CameraService>(),
                services.GetRequiredService<IlluminatorService>(),
                output,
                gainR, gainG, gainB,
                threshold, steps,
                intensityWhite
            )
        );

        // Собираем хост
        var host = builder.Build();

        // Поскольку LibCameraLogSink, который перенаправляет логи из libcamera в хост,
        // не будет создан ни одним из других сервисов, его необходимо создать превентивно
        host.Services.GetRequiredService<LibCameraLogSink>();

        // Запускаем хост
        host.Run();
    }
}