SpoutColorNode.cs

using System;
using System.Collections.Generic;
using System.Linq;
using Cysharp.Threading.Tasks;
using Klak.Spout;
using Unity.Collections;
using Warudo.Core.Attributes;
using Warudo.Core.Graphs;
using UnityEngine;
using UnityEngine.Rendering;
using Warudo.Core.Data;

namespace Oose.SpoutColorPlugin {
    [NodeType(Id = "11654261-dfe5-4f83-8498-93a9f3d0f437", Title = "Get color from Spout source", Category = "Lights")]
    public class SpoutColorNode : Node
    {
        private static readonly HashSet<SpoutColorNode> Nodes = new();
        private static readonly HashSet<SpoutReceiver> Receivers = new ();
        private static readonly HashSet<SpoutColorSource> Sources = new();

        private class SpoutColorSource : IDisposable
        {
            public readonly RenderTexture SourceTexture;
            
            private Color _finalColor;
            private readonly RenderTexture[] _intermediates;
            private readonly CommandBuffer _commandBuffer;
            private NativeArray<byte> _colors = new (FinalTextureSize*FinalTextureSize*4, Allocator.Persistent);
            private bool _requestActive;
            private int _lastRequestFrame = Time.frameCount;

            private const int FinalTextureSize = 2;

            public SpoutColorSource(RenderTexture sourceTexture)
            {
                SourceTexture = sourceTexture;
                _finalColor = Color.black;
                if (!IsValid())
                {
                    return;
                }
                _intermediates = new[]
                {
                    new RenderTexture(256,256,0),
                    new RenderTexture(64,64,0),
                    new RenderTexture(16,16,0),
                    new RenderTexture(4,4,0),
                    new RenderTexture(FinalTextureSize,FinalTextureSize,0),
                };
                _commandBuffer = new CommandBuffer();
                _commandBuffer.Blit(sourceTexture, _intermediates[0]);
                for (int i = 0; i < _intermediates.Length - 1; i++)
                {
                    _commandBuffer.Blit(_intermediates[i], _intermediates[i+1]);
                }
                _commandBuffer.RequestAsyncReadbackIntoNativeArray(ref _colors, _intermediates[^1], mipIndex: 0, TextureFormat.RGBA32, RequestCallback);
            }

            public bool IsValid()
            {
                return SourceTexture;
            }
            
            public Color GetColor()
            {
                if (!IsValid())
                {
                    return _finalColor;
                }

                if (!_requestActive && Time.frameCount >= _lastRequestFrame)
                {
                    _requestActive = true;
                    Graphics.ExecuteCommandBufferAsync(_commandBuffer, ComputeQueueType.Background);
                }
                return _finalColor;
            }
            
            void RequestCallback(AsyncGPUReadbackRequest request)
            {
                if (!IsValid() || request.hasError || !request.done || !_colors.IsCreated)
                {
                    _lastRequestFrame = Time.frameCount + 5;
                    _requestActive = false;
                    return;
                }
                var r = 0f;
                var g = 0f;
                var b = 0f;
                var pixels = _colors.Length / 4;
                for (int i = 0; i < _colors.Length / 4; i++)
                {
                    r += _colors[i * 4 + 0] / 255f;
                    g += _colors[i * 4 + 1] / 255f;
                    b += _colors[i * 4 + 2] / 255f;
                }
                _finalColor = new Color(r / pixels, g / pixels, b / pixels, 1f);
                _lastRequestFrame = Time.frameCount;
                _requestActive = false;
            }

            public void Dispose()
            {
                if (_intermediates != null)
                {
                    foreach (var rt in _intermediates)
                    {
                        rt.Release();
                    }
                }
                _commandBuffer?.Dispose();
                _colors.Dispose();
            }
        }

        [DataInput]
        [Label("Source name")]
        [AutoComplete(nameof(AutoCompleteSpoutSources), forceSelection: false)]
        public string SourceName;
        
        private async UniTask<AutoCompleteList> AutoCompleteSpoutSources() {
            return SpoutManager.GetSourceNames().Select(x => new AutoCompleteEntry {
                label = x,
                value = x
            }).ToAutoCompleteList();
        }

        [DataOutput]
        [Label("Color")]
        public Color SourceColor()
        {
            Cleanup();
            if (TryGetReceiver(SourceName, out var receiver))
            {
                if (TryGetSource(receiver.receivedTexture, out var source))
                {
                    return source.GetColor();
                }
            }
            return Color.black;
        }

        bool TryGetReceiver(string name, out SpoutReceiver receiver)
        {
            if (string.IsNullOrEmpty(name))
            {
                receiver = null;
                return false;
            }
            receiver = Receivers.FirstOrDefault(x => x.sourceName == name);
            if (receiver)
            {
                return true;
            }
            receiver = UnityEngine.Object.Instantiate(new GameObject()).AddComponent<SpoutReceiver>();
            receiver.sourceName = name;
            Debug.Log($"SpoutColorNode: Created SpoutReceiver for {name}");
            Receivers.Add(receiver);
            return true;
        }

        bool TryGetSource(RenderTexture rt, out SpoutColorSource source)
        {
            if (rt == null)
            {
                source = null;
                return false;
            }
            source = Sources.FirstOrDefault(x => x.SourceTexture == rt);
            if (source != null)
            {
                return true; 
            }
            source = new(rt);
            if (source.IsValid())
            {
                Sources.Add(source);
                Debug.Log($"SpoutColorNode: Created new SpoutColorSource instance: total {Sources.Count}");
                return true;
            }
            source.Dispose();
            return false;
        }
        
        protected override void OnCreate()
        {
            Nodes.Add(this);
        }

        protected override void OnDestroy()
        {
            if (Nodes.Contains(this))
            {
                Nodes.Remove(this);
            }
            Cleanup();
        }

        private static void Cleanup()
        {
            string[] inactiveSourceNames = Receivers.Select(x => x.sourceName).Where(x => !Nodes.Select(y => y.SourceName).Contains(x)).ToArray();
            foreach (var name in inactiveSourceNames)
            {
                foreach (var r in Receivers.Where(x => x && x.sourceName == name))
                {
                    UnityEngine.Object.Destroy(r);
                }
                Receivers.RemoveWhere(x => x && x.sourceName == name);
                Receivers.RemoveWhere(x => !x);
                Debug.Log($"SpoutColorNode: Removing SpoutReceiver for {name}");
            }

            int count = Sources.Count;
            foreach (var source in Sources)
            {
                if (!source.IsValid())
                {
                    source.Dispose();
                }
            }
            Sources.RemoveWhere(x => !x.IsValid());
            if (Sources.Count != count)
            {
                Debug.Log($"SpoutColorNode: Removed excess SpoutColorSource instances: from {count} down to {Sources.Count}");
            }
        }
        
        public static void DestroyAllSources()
        {
            foreach (var s in Sources)
            {
                s.Dispose();
            }
            Sources.Clear();
            foreach (var r in Receivers)
            {
                if (r)
                {
                    UnityEngine.Object.Destroy(r.gameObject);
                }
            }
            Receivers.Clear();
        }
    }
}