Project_2018_Frozen/Unity_2018_Frozen/Assets/KinectOpticalFlow/Grpc/KinectGrpcServer.cs

using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using Grpc.Core;
using System.Threading.Tasks;

namespace UltraCombos.Kinect2MotionTransmit
{
    public class KinectGrpcServer : MonoBehaviour
    {
        public int port = 50051;
        Server server;

        [SerializeField]
        ComputeShader shader;
        public ComputeBuffer buffer;

        [SerializeField]
        ComputeShader collectShader;

        [SerializeField]
        bool clip = true;

        class KinectBuffer
        {
            ~KinectBuffer()
            {
                Release();
            }

            public void Release()
            {
                if (position != null)
                {
                    position.Release();
                    position = null;
                }
                if (velocity != null)
                {
                    velocity.Release();
                    velocity = null;
                }
            }

            public ComputeBuffer position;
            public ComputeBuffer velocity;

            public object mutex = new object();
            public DataRequest[] data_pool = new DataRequest[3];
            public bool has_new_data = false;
            public int READY = 0;
            public int FRONT = 1;
            public int BACK = 2;

            public float fps = 30.0f;
            public float stamp = 0.0f;

            public Matrix4x4 matrix;
            public Vector3 ClipBox;
        }
        
        List<KeyValuePair<string, KinectBuffer>> kinect_buffers = new List<KeyValuePair<string, KinectBuffer>>();

        int final_count = 1;
        ComputeBuffer final_position_buffer;
        public ComputeBuffer FinalPositionBuffer { get { return final_position_buffer; } }
        ComputeBuffer final_velocity_buffer;
        public ComputeBuffer FinalVelocityBuffer { get { return final_velocity_buffer; } }
       

        Texture2D depth_texture;
        byte[] depth_data;
        //Texture2D depth_low_texture;
        //byte[] depth_low_data;
        Texture2D velocity_texture;
        byte[] velocity_data;

        const int width = 512;
        const int height = 424;


        [SerializeField]
        List<string> connectedDevices = new List<string>();
        [SerializeField]
        string debug;
        [SerializeField]
        bool debugDraw = false;
        [SerializeField]
        Shader debugShader = null;

        private void Start()
        {
            buffer = new ComputeBuffer(width * height, sizeof(float) * 4);
            Debug.LogFormat("Buffer is created by count {0} and stride {1}", buffer.count, buffer.stride);
            depth_data = new byte[buffer.count ];
            //depth_low_data = new byte[buffer.count / 2];
            velocity_data = new byte[buffer.count];

            depth_texture = new Texture2D(width, height, TextureFormat.BC5, false);
            //depth_low_texture = new Texture2D(width, height, TextureFormat.BC4, false);
            velocity_texture = new Texture2D(width, height, TextureFormat.BC5, false);
            
            final_position_buffer = new ComputeBuffer(width * height, sizeof(float) * 4);
            final_velocity_buffer = new ComputeBuffer(width * height, sizeof(float) * 4);

            StartServer();
        }

        private void Update()
        {
            uint x, y, z;
            shader.GetKernelThreadGroupSizes(0, out x, out y, out z);
            int[] args = new int[3]
            {
                (int)Mathf.CeilToInt((float)width / x),
                (int)Mathf.CeilToInt((float)height / y),
                (int)Mathf.CeilToInt((float)1 / z)
            };
            
            connectedDevices.Clear();
            debug = "";
            foreach (var buf in kinect_buffers)
            {
                connectedDevices.Add(buf.Key);
                debug += string.Format("{0} ", buf.Value.fps);

                var kinect = buf.Value;

                if (kinect.position == null)
                    kinect.position = new ComputeBuffer(width * height, sizeof(float) * 4);
                if (kinect.velocity == null)
                    kinect.velocity = new ComputeBuffer(width * height, sizeof(float) * 4);

                if (kinect.has_new_data)
                {
                    lock (kinect.mutex)
                    {
                        Swap(ref kinect.FRONT, ref kinect.READY);
                        kinect.has_new_data = false;
                    }

                    DataRequest data = kinect.data_pool[kinect.READY];

                    Vector3 pos = new Vector3(data.PositionX, data.PositionY, -data.PositionZ);
#if false
                    Quaternion q = Quaternion.Euler(-data.RotationX, -data.RotationY, data.RotationZ);
#else
                    Quaternion q = Quaternion.Euler(0, 0, data.RotationZ);
                    q *= Quaternion.Euler(0, -data.RotationY, 0);
                    q *= Quaternion.Euler(-data.RotationX, 0, 0);
#endif
                    kinect.matrix = Matrix4x4.TRS(pos, q, Vector3.one);
                    kinect.ClipBox = new Vector3(data.ClipSizeX, data.ClipSizeY, data.ClipSizeZ);

                    data.DepthHighData.CopyTo(depth_data, 0);
                    //data.DepthLowData.CopyTo(depth_low_data, 0);
                    data.VelocityData.CopyTo(velocity_data, 0);

                    depth_texture.LoadRawTextureData(depth_data);
                    depth_texture.Apply();
                    //depth_low_texture.LoadRawTextureData(depth_low_data);
                    //depth_low_texture.Apply();
                    velocity_texture.LoadRawTextureData(velocity_data);
                    velocity_texture.Apply();

                    shader.SetTexture(0, "depth_texture", depth_texture);
                    //shader.SetTexture(0, "depth_low_texture", depth_low_texture);
                    shader.SetTexture(0, "velocity_texture", velocity_texture);
                    shader.SetBuffer(0, "kinect_motion_buffer", buffer);
                    shader.SetBuffer(0, "position_buffer", kinect.position);
                    shader.SetBuffer(0, "velocity_buffer", kinect.velocity);
                    shader.SetMatrix("kinect_matrix", kinect.matrix);
                    shader.SetVector("clip_box", kinect.ClipBox);
                    shader.SetFloat("clip", clip ? 1.0f : 0.0f);
                    shader.Dispatch(0, args[0], args[1], args[2]);

                    kinect.fps = Mathf.Lerp(kinect.fps, 1.0f / (Time.time - kinect.stamp), 0.02f);
                    kinect.stamp = Time.time;
                }
            }
            
        }

        private void FixedUpdate()
        {
            uint x, y, z;
            collectShader.GetKernelThreadGroupSizes(0, out x, out y, out z);
            int[] args = new int[3]
            {
                (int)Mathf.CeilToInt((float)(width * height) / x),
                (int)Mathf.CeilToInt((float)1 / y),
                (int)Mathf.CeilToInt((float)1 / z)
            };
            
            if (final_count < kinect_buffers.Count)
            {
                final_count = kinect_buffers.Count;
                int buf_count = final_count * width * height;
                final_position_buffer.Release();
                final_position_buffer = new ComputeBuffer(buf_count, sizeof(float) * 4);
                final_velocity_buffer.Release();
                final_velocity_buffer = new ComputeBuffer(buf_count, sizeof(float) * 4);
            }

            int count = 0;
            foreach (var buf in kinect_buffers)
            {
                if (buf.Value.position == null || buf.Value.velocity == null)
                    continue;
                collectShader.SetInt("count", count++);
                collectShader.SetBuffer(0, "position_buffer", buf.Value.position);
                collectShader.SetBuffer(0, "velocity_buffer", buf.Value.velocity);
                collectShader.SetBuffer(0, "FinalPositionBuffer", final_position_buffer);
                collectShader.SetBuffer(0, "FinalVelocityBuffer", final_velocity_buffer);
                collectShader.Dispatch(0, args[0], args[1], args[2]);
            }
                
        }

        private void OnDestroy()
        {
            server.ShutdownAsync().Wait();

            if (buffer != null)
            {
                buffer.Release();
                buffer = null;
            }

            if (final_position_buffer != null)
            {
                final_position_buffer.Release();
                final_position_buffer = null;
            }

            if (final_velocity_buffer != null)
            {
                final_velocity_buffer.Release();
                final_velocity_buffer = null;
            }

            foreach (var buf in kinect_buffers)
            {
                buf.Value.Release();
            }
        }

        private void StartServer()
        {
            try
            {
                var options = new List<ChannelOption> { new ChannelOption(ChannelOptions.MaxReceiveMessageLength, int.MaxValue) };
                server = new Server(options)
                {
                    Services = { Resource.BindService(new ResourceImpl(OnMessageReceived)) },
                    Ports = { new ServerPort("0.0.0.0", port, ServerCredentials.Insecure) }
                };

                server.Start();
                Debug.Log("Start Server");
            }
            catch (System.Exception e)
            {
                Debug.LogError(e.Message);
            }
        }

        private void OnMessageReceived(DataRequest data)
        {
            //Debug.Log(data.Name);
            //Debug.LogFormat("message.Length: {0}", data.Content.Length);
            
            KinectBuffer kinect = null;
            foreach (var buf in kinect_buffers)
            {
                if (data.Name.Equals(buf.Key))
                {
                    kinect = buf.Value;
                } 
            }

            if (kinect == null)
            {
                kinect = new KinectBuffer();
                for (int i = 0; i < 3; i++)
                    kinect.data_pool[i] = new DataRequest();
                kinect_buffers.Add(new KeyValuePair<string, KinectBuffer>(data.Name, kinect));
            }

            kinect.data_pool[kinect.BACK].MergeFrom(data);
            lock (kinect.mutex)
            {
                Swap(ref kinect.BACK, ref kinect.FRONT);
                kinect.has_new_data = true;
            }
        }

        private static void Swap<T>(ref T a, ref T b)
        {
            T tmp = a;
            a = b;
            b = tmp;
        }

        private void OnDrawGizmosSelected()
        {
            var col = Gizmos.color;
            float fov = 60.0f;
            float aspect = 512.0f / 424.0f;
            int count = 0;
            foreach (var buf in kinect_buffers)
            {
                var kinect = buf.Value;
                Gizmos.color = Color.HSVToRGB(count++ / 5.0f, 0.8f, 0.8f);
                var mat = Gizmos.matrix;
                Gizmos.matrix = transform.localToWorldMatrix;
                Gizmos.matrix = Matrix4x4.identity;
                Gizmos.DrawWireCube(new Vector3(0, kinect.ClipBox.y * 0.5f, 0), kinect.ClipBox);
                Gizmos.matrix *= kinect.matrix;
                Gizmos.DrawFrustum(Vector3.zero, fov, 10.0f, 0.0f, aspect);
                Gizmos.matrix = mat;
            }

            Gizmos.color = col;
        }

        Material material = null;

        private void OnRenderObject()
        {
            if (debugShader == null)
                return;

            if (debugDraw == false)
                return;

            if (material == null)
                material = new Material(debugShader);

            material.SetPass(0);
            material.SetBuffer("ssbo", final_position_buffer);

            Graphics.DrawProcedural(MeshTopology.Points, final_position_buffer.count, 1);

        }
    }

    internal class ResourceImpl : Resource.ResourceBase
    {
        public ResourceImpl(MessageDelegate func = null)
        {
            if (func != null)
                onMessageReceived += func;
        }

        public override Task<DataReply> SendData(DataRequest request, ServerCallContext context)
        {
            onMessageReceived.Invoke(request);
            string res = string.Format("OK: DH({0}), DL({1}), V({2})", request.DepthHighData.Length, request.DepthLowData.Length, request.VelocityData.Length);
            //Debug.Log(res);
            return Task.FromResult(new DataReply { Result = res });
        }

        public delegate void MessageDelegate(DataRequest data);
        public MessageDelegate onMessageReceived;
    }
}