Project_2018_Frozen/Unity_2018_Frozen/Assets/KinectOpticalFlow/KinectOpticalFlow.cs

using System.Collections;
using System.Collections.Generic;
using UnityEngine;

namespace UltraCombos
{
    public class KinectOpticalFlow : MonoBehaviour
    {
        const int width = 512;
        const int height = 424;
        const int POSITION = 0;
        const int VELOCITY = 1;

        public Kinect2MotionTransmit.KinectGrpcServer server;

        [SerializeField]
        List<StructuredBuffer> kinectBuffers = new List<StructuredBuffer>();
        public ComputeBuffer PositionBuffer { get { return kinectBuffers[POSITION].obj; } }
        public ComputeBuffer VelocityBuffer { get { return kinectBuffers[VELOCITY].obj; } }

        [SerializeField]
        ComputeShader updateShader;
        [SerializeField]
        ComputeShader applyShader;
        ComputeBuffer clear_grid_buffer_args_buffer = null;
        ComputeBuffer args_buffer = null;

        public bool flipX = true;
        public Vector3 clipMin = Vector3.one * -10;
        public Vector3 clipMax = Vector3.one * 10;


        [SerializeField]
        FluidSim3DProject.SmokeFluidSim fluidSimulation;
        ComputeBuffer kinect_grid_buffer;
        public ComputeBuffer KinectGridBuffer { get { return kinect_grid_buffer; } }
        int grid_width, grid_height, grid_depth;



        [SerializeField, Header("Debug")]
        bool debug = false;
        [SerializeField]
        bool colorize = false;
        [SerializeField]
        Material debugMaterial;

        ComputeBuffer home_position_buffer;
        public ComputeBuffer HomePositionBuffer { get { return home_position_buffer; } }
        ComputeBuffer home_velocity_buffer;
        public ComputeBuffer HomeVelocityBuffer { get { return home_velocity_buffer; } }
        ComputeBuffer home_count_buffer;
        public ComputeBuffer HomeCountBuffer { get { return home_count_buffer; } }

        


        private void Start()
        {
            foreach (var buf in kinectBuffers)
                buf.Allocate(width * height, sizeof(float) * 4);

            {
                grid_width = fluidSimulation.m_width;
                grid_height = fluidSimulation.m_height;
                grid_depth = fluidSimulation.m_depth;
                int grid_count = grid_width * grid_height * grid_depth;
                kinect_grid_buffer = new ComputeBuffer(grid_count, sizeof(float) * 4);
            }

            {
                home_position_buffer = new ComputeBuffer(512 * 424, sizeof(float) * 4, ComputeBufferType.Append);
                home_velocity_buffer = new ComputeBuffer(512 * 424, sizeof(float) * 4, ComputeBufferType.Append);
                home_count_buffer = new ComputeBuffer(1, sizeof(int), ComputeBufferType.Raw);
            }

            
        }
        
        private void FixedUpdate()
        {
            if (args_buffer == null)
            {
                args_buffer = updateShader.CreateIndirectComputeArgsBuffer(width * height, 1, 1);
            }            
            
            if (server != null)
            {
                if (home_position_buffer.count != server.FinalPositionBuffer.count)
                {
                    home_position_buffer.Release();
                    home_position_buffer = new ComputeBuffer(server.FinalPositionBuffer.count,
                        server.FinalPositionBuffer.stride, ComputeBufferType.Append);
                }

                if (home_velocity_buffer.count != server.FinalVelocityBuffer.count)
                {
                    home_velocity_buffer.Release();
                    home_velocity_buffer = new ComputeBuffer(server.FinalVelocityBuffer.count,
                        server.FinalVelocityBuffer.stride, ComputeBufferType.Append);
                }

                foreach (var buf in kinectBuffers)
                {
                    if (buf.count != server.FinalPositionBuffer.count)
                    {
                        buf.Release();
                        buf.Allocate(server.FinalVelocityBuffer.count, server.FinalVelocityBuffer.stride);
                    }
                }
            }
            
            // update 
            {
                var m = transform.localToWorldMatrix;
                float[] mat_floats = new float[16]
                { m.m00, m.m01, m.m02, m.m03,
                  m.m10, m.m11, m.m12, m.m13,
                  m.m20, m.m21, m.m22, m.m23,
                  m.m30, m.m31, m.m32, m.m33};

                updateShader.SetVector("ClipMin", clipMin);
                updateShader.SetVector("ClipMax", clipMax);
                updateShader.SetFloats("ModelMatrix", mat_floats);
                updateShader.SetVector("FluidSize", fluidSimulation.GridSize);
                updateShader.SetVector("FluidRoot", fluidSimulation.GridRoot);
                updateShader.SetVector("FluidDim", fluidSimulation.GridDim);
                updateShader.SetInt("flipX", flipX ? 1 : 0);
                
                foreach (var buf in kinectBuffers)
                    updateShader.SetBuffer(0, buf.bufferName, buf.obj);
                updateShader.SetBuffer(0, "KinectGridBuffer", kinect_grid_buffer);
                updateShader.SetBuffer(0, "HomePositionBuffer", home_position_buffer);
                updateShader.SetBuffer(0, "HomeVelocityBuffer", home_velocity_buffer);

                if (server != null)
                {
                    updateShader.SetBuffer(0, "FinalPositionBuffer", server.FinalPositionBuffer);
                    updateShader.SetBuffer(0, "FinalVelocityBuffer", server.FinalVelocityBuffer);

                    uint x, y, z;
                    updateShader.GetKernelThreadGroupSizes(0, out x, out y, out z);
                    int tgx = Mathf.CeilToInt((float)server.FinalPositionBuffer.count / x);

                    updateShader.Dispatch(0, tgx, 1, 1);
                }
                else
                {
                    updateShader.DispatchIndirect(0, args_buffer);
                }
            }

            {

                ComputeBuffer.CopyCount(home_position_buffer, home_count_buffer, 0);
#if false
                int home_count;
                var data = new int[home_count_buffer.count];
                home_count_buffer.GetData(data);
                home_count = data[0];
                Debug.LogFormat("home_count: {0}", home_count);
#endif
                home_position_buffer.SetCounterValue(0);
                home_velocity_buffer.SetCounterValue(0);
            }

            if (clear_grid_buffer_args_buffer == null)
            {
                Vector3Int dim = Vector3Int.FloorToInt(fluidSimulation.GridDim);
                clear_grid_buffer_args_buffer = applyShader.CreateIndirectComputeArgsBuffer(dim.x, dim.y, dim.z);
            }

            {
                applyShader.SetVector("FluidDim", fluidSimulation.GridDim);
                applyShader.SetBuffer(0, "VelocityBuffer", fluidSimulation.VelocityBuffer);
                applyShader.SetBuffer(0, "TemperatureBuffer", fluidSimulation.TemperatureBuffer);
                applyShader.SetBuffer(0, "ObstacleBuffer", fluidSimulation.ObstacleBuffer);
                applyShader.SetBuffer(0, "DensityBuffer", fluidSimulation.DensityBuffer);
                applyShader.SetFloat("timeStep", Time.fixedDeltaTime);
                applyShader.SetFloat("temperatureAmount", fluidSimulation.m_temperatureAmount);
                applyShader.SetFloat("densityAmount", fluidSimulation.m_densityAmount);
                applyShader.SetBuffer(0, "KinectGridBuffer", kinect_grid_buffer);
                applyShader.DispatchIndirect(0, clear_grid_buffer_args_buffer);
            }
        }

        private bool IsTextureValid(Texture2D tex)
        {
            bool res = false;
            if (tex)
            {
                if (tex.width == 512 && tex.height == 424)
                    res = true;
            }
            return res;
        }

        private void OnRenderObject()
        {
            if (debug == false)
                return;
            if (Utilities.SkipCamera(Camera.current.name))
                return;

            debugMaterial.SetPass(0);
            debugMaterial.SetInt("colorize", colorize ? 1 : 0);
            foreach (var buf in kinectBuffers)
                debugMaterial.SetBuffer(buf.bufferName, buf.obj);
            Graphics.DrawProcedural(MeshTopology.Points, PositionBuffer.count);
        }

        private void OnDestroy()
        {
            clear_grid_buffer_args_buffer.Release();
            args_buffer.Release();
            KinectGridBuffer.Release();

            home_position_buffer.Release();
            home_velocity_buffer.Release();
            home_count_buffer.Release();
        }

    }

}