LlamaCasty/Services/BackgroundEffects/OnnxPersonSegmenter.cs

using System.IO;
using Microsoft.ML.OnnxRuntime;
using ytLive.Helpers;

namespace ytLive.Services.BackgroundEffects;

/// <summary>
/// The default <see cref="IPersonSegmenter"/>: Google's MediaPipe Selfie Segmenter
/// (landscape, Apache-2.0) converted to ONNX by <c>tools/selfie-seg/</c> and shipped as an
/// embedded resource — nothing is downloaded at runtime and no GPU is required.
///
/// Contract: input [1,144,256,3] RGB in [0,1] (NHWC), output [1,144,256,1] person
/// probability. The whole frame is resampled to 256×144 (stretched, not cropped), so the
/// mask maps back to the frame by normalized coordinates.
///
/// Never throws from <see cref="Segment"/>: a model that fails to load or run logs once
/// and returns null, and the caller passes the webcam frame through.
/// </summary>
public sealed class OnnxPersonSegmenter : IPersonSegmenter, IDisposable
{
    public const int ModelWidth = 256;
    public const int ModelHeight = 144;
    internal const string ModelResourceName = "ytLive.Assets.Models.selfie_segmenter_landscape.onnx";

    private static readonly long[] InputShape = { 1, ModelHeight, ModelWidth, 3 };

    private readonly object _gate = new();
    private readonly float[] _input = new float[ModelWidth * ModelHeight * 3];
    private readonly RunOptions _runOptions = new();
    private InferenceSession? _session;
    private bool _failed;
    private bool _disposed;

    public SegmentationMask? Segment(VideoFrame frame)
    {
        lock (_gate)
        {
            if (_failed || _disposed) return null;
            try
            {
                _session ??= CreateSession();
                FillInput(frame, _input);
                using var input = OrtValue.CreateTensorValueFromMemory(_input, InputShape);
                using var outputs = _session.Run(_runOptions, new[] { "input" }, new[] { input }, new[] { "mask" });
                var mask = outputs[0].GetTensorDataAsSpan<float>().ToArray();
                return new SegmentationMask(ModelWidth, ModelHeight, mask);
            }
            catch (Exception ex)
            {
                _failed = true;
                AppLog.Write($"OnnxPersonSegmenter: background effects disabled — {ex.GetType().Name}: {ex.Message}");
                return null;
            }
        }
    }

    public void Dispose()
    {
        lock (_gate)
        {
            _disposed = true;
            _session?.Dispose();
            _session = null;
            _runOptions.Dispose();
        }
    }

    private static InferenceSession CreateSession()
    {
        using var stream = typeof(OnnxPersonSegmenter).Assembly.GetManifestResourceStream(ModelResourceName)
            ?? throw new InvalidOperationException($"embedded model '{ModelResourceName}' is missing");
        using var buffer = new MemoryStream();
        stream.CopyTo(buffer);
        // Two intra-op threads: the model is tiny, and the webcam must not starve the
        // compositor/encoder threads (obs-backgroundremoval reports 2 as the sweet spot).
        using var options = new SessionOptions
        {
            IntraOpNumThreads = 2,
            InterOpNumThreads = 1,
            GraphOptimizationLevel = GraphOptimizationLevel.ORT_ENABLE_ALL,
        };
        return new InferenceSession(buffer.ToArray(), options);
    }

    /// <summary>BGRA frame → model input: nearest-sample stretch to 256×144, RGB order,
    /// values /255. Nearest is enough — the model input is ~5× smaller than a 720p frame
    /// and the mask is the only thing that comes back.</summary>
    internal static void FillInput(VideoFrame frame, float[] input)
    {
        var src = frame.BgraPixels;
        var stride = frame.Stride;
        var o = 0;
        for (var y = 0; y < ModelHeight; y++)
        {
            var sy = Math.Min(frame.Height - 1, (int)((y + 0.5) * frame.Height / ModelHeight));
            var row = sy * stride;
            for (var x = 0; x < ModelWidth; x++)
            {
                var sx = Math.Min(frame.Width - 1, (int)((x + 0.5) * frame.Width / ModelWidth));
                var p = row + sx * 4;
                input[o++] = src[p + 2] / 255f;
                input[o++] = src[p + 1] / 255f;
                input[o++] = src[p] / 255f;
            }
        }
    }
}