LlamaCasty/Services/ChromaKeyDetector.cs

namespace ytLive.Services;

/// <summary>What a clip's background looks like, as far as keying goes.</summary>
public enum ChromaScreen
{
    /// <summary>No usable backdrop — say nothing to the creator.</summary>
    None,
    Green,
    Blue,
}

/// <summary>
/// TASK 64: "this clip was shot against a green screen". <c>Coverage</c> × <c>Tightness</c>:
/// how much of the frame is the colour, and how consistent that colour is. A painted screen
/// scores near 1; a lawn or a football pitch scores low because the greens disagree.
/// </summary>
public readonly record struct ChromaKeySuggestion(
    ChromaScreen Screen,
    string? ColorHex,
    double Confidence)
{
    /// <summary>The screen is uniform enough to key without the creator touching a slider.</summary>
    public bool IsHighConfidence => Screen != ChromaScreen.None && Confidence >= HighConfidenceFloor;

    public const double HighConfidenceFloor = 0.85;
    public const double SuggestFloor = 0.65;
}

/// <summary>
/// TASK 64: finds a green or blue screen in a decoded frame — the "clueless creator" half of
/// the feature. Pure and allocation-light (it only samples an outer ring and a coarse grid), so
/// it runs on the probe frame without touching ffmpeg twice.
/// <para><b>It never decides anything.</b> A suggestion is offered to the creator; auto-apply is
/// an explicit setting and only fires at <see cref="ChromaKeySuggestion.IsHighConfidence"/>.
/// Wrongly keying a lawn is destructive, so "None" is a normal, expected answer.</para>
/// </summary>
public static class ChromaKeyDetector
{
    // Hue windows in degrees. Green sits 75–165 (a real screen reads ~120–145), blue 185–265.
    private const double GreenHueFrom = 75, GreenHueTo = 165;
    private const double BlueHueFrom = 185, BlueHueTo = 265;

    // A screen is saturated and not dark. These also exclude letterbox bars (near-black,
    // no saturation) and blown-out highlights from the denominator, so black bars do not
    // count against an otherwise perfect screen.
    private const double MinSaturation = 0.35;
    private const double MinValue = 0.15;
    private const double MaxValue = 0.98;

    // The screen's own pixels must agree with each other: a painted surface is flat, a lawn
    // is not. This is what keeps "green field" from being keyed.
    private const double HueTightness = 12.0;

    /// <summary>Samples an outer ring (the subject sits in the middle, the screen fills the
    /// edges) and reports what it saw.</summary>
    public static ChromaKeySuggestion Detect(VideoFrame frame)
    {
        if (frame.Width < 8 || frame.Height < 8) return Empty;

        var ring = Math.Max(2, (int)Math.Round(Math.Min(frame.Width, frame.Height) * 0.15));
        var step = Math.Max(1, Math.Min(frame.Width, frame.Height) / 96);

        var hues = new List<double>();
        var sameHue = 0;
        double hueSum = 0;
        long rSum = 0, gSum = 0, bSum = 0;
        var considered = 0;
        var green = 0;
        var blue = 0;

        for (var y = 0; y < frame.Height; y += step)
        {
            var inVerticalRing = y < ring || y >= frame.Height - ring;
            for (var x = 0; x < frame.Width; x += step)
            {
                if (!inVerticalRing && x >= ring && x < frame.Width - ring) continue;

                var i = y * frame.Stride + x * 4;
                var b = frame.BgraPixels[i];
                var g = frame.BgraPixels[i + 1];
                var r = frame.BgraPixels[i + 2];
                if (frame.BgraPixels[i + 3] < 8) continue; // transparent: not the backdrop

                var (hue, saturation, value) = ToHsv(r, g, b);
                if (saturation < MinSaturation || value < MinValue || value > MaxValue) continue;

                considered++;
                if (hue is >= GreenHueFrom and <= GreenHueTo) green++;
                else if (hue is >= BlueHueFrom and <= BlueHueTo) blue++;
                else continue;

                hues.Add(hue);
                hueSum += hue;
                rSum += r;
                gSum += g;
                bSum += b;
            }
        }

        if (considered < 64) return Empty;

        var screen = green >= blue
            ? (green == 0 ? ChromaScreen.None : ChromaScreen.Green)
            : ChromaScreen.Blue;
        if (screen == ChromaScreen.None) return Empty;

        var accepted = hues.Count;
        var coverage = (double)accepted / considered;

        // Tint stability: how many of the accepted pixels sit within ±HueTightness of the
        // median hue. A painted screen ≈ 1.0; grass, foliage and sunlight ≈ far less.
        hues.Sort();
        var median = hues[hues.Count / 2];
        foreach (var hue in hues)
        {
            var delta = Math.Abs(hue - median);
            if (delta > 180) delta = 360 - delta;
            if (delta <= HueTightness) sameHue++;
        }
        var tightness = (double)sameHue / accepted;

        var confidence = coverage * tightness;
        if (confidence < ChromaKeySuggestion.SuggestFloor) return Empty;

        // The keyed colour is the screen as it was actually shot, not a canonical green:
        // similarity is a distance, so the median of the sampled pixels keys far better.
        var keyed = Math.Max(1, accepted);
        var color = $"#{rSum / keyed:X2}{gSum / keyed:X2}{bSum / keyed:X2}";
        return new ChromaKeySuggestion(screen, color, confidence);
    }

    private static readonly ChromaKeySuggestion Empty = new(ChromaScreen.None, null, 0);

    private static (double Hue, double Saturation, double Value) ToHsv(byte r, byte g, byte b)
    {
        var rf = r / 255.0;
        var gf = g / 255.0;
        var bf = b / 255.0;
        var max = Math.Max(rf, Math.Max(gf, bf));
        var min = Math.Min(rf, Math.Min(gf, bf));
        var delta = max - min;

        double hue = 0;
        if (delta > 1e-6)
        {
            if (max == rf) hue = 60 * (((gf - bf) / delta) % 6);
            else if (max == gf) hue = 60 * (((bf - rf) / delta) + 2);
            else hue = 60 * (((rf - gf) / delta) + 4);
        }
        if (hue < 0) hue += 360;

        var saturation = max <= 0 ? 0 : delta / max;
        return (hue, saturation, max);
    }
}