HawkeyeVision/Hawkeye.VisionBuilder.UI.RecipeYoloPatcher/MainForm.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Text;
using System.Windows.Forms;
using Compunet.YoloV8;
using Hawkeye.VisionBuilder.Workflow;
using Hawkeye.VisionBuilder.Workflow.Operations.AI;

namespace Hawkeye.VisionBuilder.UI.RecipeYoloPatcher;

public partial class MainForm : Form
{
    private sealed class ClassEntry
    {
        public int Id { get; init; }
        public string Name { get; init; } = "";
        public override string ToString() => $"{Id}: {Name}";
    }

    public MainForm()
    {
        InitializeComponent();
    }

    private void BtnBrowseModel_Click(object? sender, EventArgs e)
    {
        using var dialog = new OpenFileDialog
        {
            Filter = "ONNX models (*.onnx)|*.onnx",
            Title = "Select Yolo model"
        };
        var defaultDir = Path.GetFullPath(@"..\Data\Models");
        if (Directory.Exists(defaultDir)) dialog.InitialDirectory = defaultDir;
        if (dialog.ShowDialog(this) == DialogResult.OK)
        {
            txtModelPath.Text = dialog.FileName;
            btnLoadClasses.Enabled = true;
            lstClasses.Items.Clear();
        }
    }

    private void BtnLoadClasses_Click(object? sender, EventArgs e)
    {
        if (string.IsNullOrWhiteSpace(txtModelPath.Text) || !File.Exists(txtModelPath.Text))
        {
            MessageBox.Show(this, "Pick a valid .onnx model first.", "Recipe Yolo Patcher",
                MessageBoxButtons.OK, MessageBoxIcon.Warning);
            return;
        }
        lstClasses.Items.Clear();
        Cursor = Cursors.WaitCursor;
        try
        {
            using var predictor = YoloV8Predictor.Create(txtModelPath.Text);
            foreach (var cls in predictor.Metadata.Names)
            {
                lstClasses.Items.Add(new ClassEntry { Id = cls.Id, Name = cls.Name });
            }
            Log($"Loaded {lstClasses.Items.Count} class(es) from {Path.GetFileName(txtModelPath.Text)}.");
        }
        catch (Exception ex)
        {
            MessageBox.Show(this, $"Failed to load model:\r\n{ex.Message}", "Recipe Yolo Patcher",
                MessageBoxButtons.OK, MessageBoxIcon.Error);
        }
        finally
        {
            Cursor = Cursors.Default;
        }
    }

    private void BtnAddFiles_Click(object? sender, EventArgs e)
    {
        using var dialog = new OpenFileDialog
        {
            Filter = "Recipes (*.hrcp;*.jhrcp)|*.hrcp;*.jhrcp",
            Multiselect = true,
            Title = "Select recipe files"
        };
        var defaultDir = Path.GetFullPath(@"..\Data\Recipes");
        if (Directory.Exists(defaultDir)) dialog.InitialDirectory = defaultDir;
        if (dialog.ShowDialog(this) == DialogResult.OK)
        {
            foreach (var f in dialog.FileNames)
            {
                if (!lstFiles.Items.Contains(f)) lstFiles.Items.Add(f);
            }
        }
    }

    private void BtnRemoveFiles_Click(object? sender, EventArgs e)
    {
        var selected = lstFiles.SelectedItems.Cast<object>().ToList();
        foreach (var item in selected) lstFiles.Items.Remove(item);
    }

    private void BtnClearFiles_Click(object? sender, EventArgs e) => lstFiles.Items.Clear();

    private void BtnApply_Click(object? sender, EventArgs e)
    {
        var checkedClasses = lstClasses.CheckedItems.Cast<ClassEntry>().ToList();
        if (checkedClasses.Count == 0)
        {
            MessageBox.Show(this, "Check at least one class to apply.", "Recipe Yolo Patcher",
                MessageBoxButtons.OK, MessageBoxIcon.Warning);
            return;
        }
        var files = lstFiles.Items.Cast<string>().ToList();
        if (files.Count == 0)
        {
            MessageBox.Show(this, "Add at least one recipe file.", "Recipe Yolo Patcher",
                MessageBoxButtons.OK, MessageBoxIcon.Warning);
            return;
        }

        var bounds = new PatchBounds(
            (int)numMinArea.Value, (int)numMaxArea.Value,
            (int)numMinWidth.Value, (int)numMaxWidth.Value,
            (int)numMinHeight.Value, (int)numMaxHeight.Value);

        Cursor = Cursors.WaitCursor;
        btnApply.Enabled = false;
        try
        {
            foreach (var file in files)
            {
                try
                {
                    var (added, updated) = PatchRecipe(file, checkedClasses, bounds);
                    Log($"OK: {Path.GetFileName(file)} — added {added}, updated {updated}");
                }
                catch (Exception ex)
                {
                    Log($"FAIL: {Path.GetFileName(file)} — {ex.Message}");
                }
            }
            Log("Done.");
        }
        finally
        {
            Cursor = Cursors.Default;
            btnApply.Enabled = true;
        }
    }

    private void BtnStandardize_Click(object? sender, EventArgs e)
    {
        var files = lstFiles.Items.Cast<string>().ToList();
        if (files.Count == 0)
        {
            MessageBox.Show(this, "Add at least one recipe file.", "Recipe Yolo Patcher",
                MessageBoxButtons.OK, MessageBoxIcon.Warning);
            return;
        }

        Cursor = Cursors.WaitCursor;
        btnApply.Enabled = false;
        btnStandardize.Enabled = false;
        try
        {
            foreach (var file in files)
            {
                try
                {
                    var renamed = StandardizeLabels(file);
                    Log($"OK: {Path.GetFileName(file)} — renamed {renamed} label(s)");
                }
                catch (Exception ex)
                {
                    Log($"FAIL: {Path.GetFileName(file)} — {ex.Message}");
                }
            }
            Log("Done.");
        }
        finally
        {
            Cursor = Cursors.Default;
            btnApply.Enabled = true;
            btnStandardize.Enabled = true;
        }
    }

    private static int StandardizeLabels(string path)
    {
        var (wf, ext) = LoadRecipe(path);

        int renamed = 0;
        foreach (var op in wf.Operations.OfType<YoloPickDetected>())
        {
            var newLabel = ToPascalCase(op.Label);
            if (newLabel != op.Label)
            {
                op.Label = newLabel;
                renamed++;
            }
        }

        if (renamed > 0) SaveRecipe(wf, path, ext);
        return renamed;
    }

    private static string ToPascalCase(string? label)
    {
        if (string.IsNullOrWhiteSpace(label)) return label ?? string.Empty;
        var parts = label.Split((char[]?)null, StringSplitOptions.RemoveEmptyEntries);
        var sb = new StringBuilder(label.Length);
        foreach (var p in parts)
        {
            sb.Append(char.ToUpperInvariant(p[0]));
            if (p.Length > 1) sb.Append(p, 1, p.Length - 1);
        }
        return sb.ToString();
    }

    private static (WorkflowList wf, string ext) LoadRecipe(string path)
    {
        var ext = Path.GetExtension(path).ToLowerInvariant();
        WorkflowList wf;
        if (ext == ".hrcp")
        {
            wf = new WorkflowList();
            using var fs = File.OpenRead(path);
            using var br = new BinaryReader(fs);
            wf.Load(br, new OperationDiscoveryService(wf));
        }
        else if (ext == ".jhrcp")
        {
            wf = WorkflowList.LoadJSONFromFile(path);
        }
        else
        {
            throw new InvalidOperationException($"Unsupported extension '{ext}'.");
        }
        return (wf, ext);
    }

    private static void SaveRecipe(WorkflowList wf, string path, string ext)
    {
        if (ext == ".hrcp")
        {
            using var fs = File.Create(path);
            using var bw = new BinaryWriter(fs);
            wf.Save(bw);
        }
        else
        {
            wf.SaveJSON(path);
        }
    }

    private readonly record struct PatchBounds(int MinArea, int MaxArea, int MinWidth, int MaxWidth, int MinHeight, int MaxHeight);

    private static (int added, int updated) PatchRecipe(string path, IReadOnlyList<ClassEntry> classes, PatchBounds b)
    {
        var (wf, ext) = LoadRecipe(path);

        int anchor = -1;
        for (int i = 0; i < wf.Operations.Count; i++)
        {
            if (wf.Operations[i] is YoloPickDetected) anchor = i;
        }
        if (anchor < 0) anchor = wf.Operations.Count - 1;

        var nameMap = wf.Operations.OfType<YoloPickDetected>()
            .Where(op => !string.IsNullOrEmpty(op.Label))
            .GroupBy(op => op.Label)
            .ToDictionary(g => g.Key, g => g.First());

        int added = 0, updated = 0;
        foreach (var cls in classes)
        {
            if (nameMap.TryGetValue(cls.Name, out var existing))
            {
                existing.MinArea = b.MinArea;
                existing.MaxArea = b.MaxArea;
                existing.MinWidth = b.MinWidth;
                existing.MaxWidth = b.MaxWidth;
                existing.MinHeight = b.MinHeight;
                existing.MaxHeight = b.MaxHeight;
                updated++;
            }
            else
            {
                var op = new YoloPickDetected
                {
                    Label = cls.Name,
                    TypeId = cls.Id + 1,
                    SlotName = "Image",
                    MinArea = b.MinArea,
                    MaxArea = b.MaxArea,
                    MinWidth = b.MinWidth,
                    MaxWidth = b.MaxWidth,
                    MinHeight = b.MinHeight,
                    MaxHeight = b.MaxHeight
                };
                wf.Operations.Insert(++anchor, op);
                nameMap[cls.Name] = op;
                added++;
            }
        }

        SaveRecipe(wf, path, ext);
        return (added, updated);
    }

    private void Log(string msg)
    {
        txtLog.AppendText(msg + Environment.NewLine);
    }
}