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;

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

        [SerializeField]
        ComputeShader updateShader;
        [SerializeField]
        ComputeShader clearShader;
        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;
        [Range(0, 5)]
        public float kinectAmount = 0;

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

        public List<RenderTexture> renderTextures;

        private void Start()
        {
            List<string> spout_names = new List<string>() { "Position", "Velocity" };

            foreach (var name in spout_names)
            {
                var buf = new ComputeBuffer(width * height, sizeof(float) * 4);
                kinectBuffers.Add(buf);
            }

            {
                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);
            }
        }

        bool is_initialized = false;

        private void FixedUpdate()
        {
            if (is_initialized == false)
            {
                is_initialized = true;

                List<string> spout_names = new List<string>() { "Position", "Velocity" };
                int count = 0;
                foreach (var name in spout_names)
                {
                    updateShader.SetBuffer(0, string.Format("{0}Buffer", name), kinectBuffers[count]);
                    debugMaterial.SetBuffer(string.Format("{0}Buffer", name), kinectBuffers[count]);

                    count++;
                }

                {
                    updateShader.SetBuffer(0, "KinectGridBuffer", kinect_grid_buffer);
                    clearShader.SetBuffer(0, "KinectGridBuffer", kinect_grid_buffer);
                }
            }

            if (args_buffer == null)
            {
#if false
                int work_group_size = 512;
                int num_groups = (width * height + work_group_size - 1) / work_group_size;
                uint[] args = new uint[3] { (uint)num_groups, 1, 1 };
                args_buffer = new ComputeBuffer(1, args.Length * sizeof(uint), ComputeBufferType.IndirectArguments);
                args_buffer.SetData(args);
#else
                args_buffer = updateShader.CreateIndirectComputeArgsBuffer(width * height, 1, 1);
#endif
            }

            // reset kinect buffer counter
            {
                foreach (var buf in kinectBuffers)
                {
                    buf.SetCounterValue(0);
                }
            }

            // 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);

                updateShader.SetTexture(0, "PositionTexture", renderTextures[POSITION]);
                updateShader.SetTexture(0, "VelocityTexture", renderTextures[VELOCITY]);
                updateShader.DispatchIndirect(0, args_buffer);
            }

            if (clear_grid_buffer_args_buffer == null)
            {
#if false
                uint work_group_size = 8;
                Vector3 dim = fluidSimulation.GridDim;
                uint[] args = new uint[3] { (uint)dim.x / work_group_size, (uint)dim.y / work_group_size, (uint)dim.z / work_group_size };
                clear_grid_buffer_args_buffer = new ComputeBuffer(1, args.Length * sizeof(uint), ComputeBufferType.IndirectArguments);
                clear_grid_buffer_args_buffer.SetData(args);
#else
                Vector3Int dim = Vector3Int.FloorToInt(fluidSimulation.GridDim);
                clear_grid_buffer_args_buffer = clearShader.CreateIndirectComputeArgsBuffer(dim.x, dim.y, dim.z);
#endif
            }

            {
                clearShader.SetVector("FluidDim", fluidSimulation.GridDim);
                clearShader.SetBuffer(0, "VelocityBuffer", fluidSimulation.VelocityBuffer);
                clearShader.SetBuffer(0, "TemperatureBuffer", fluidSimulation.TemperatureBuffer);
                clearShader.SetBuffer(0, "ObstacleBuffer", fluidSimulation.ObstacleBuffer);
                clearShader.SetBuffer(0, "DensityBuffer", fluidSimulation.DensityBuffer);
                clearShader.SetFloat("timeStep", Time.fixedDeltaTime);
                clearShader.SetFloat("temperatureAmount", fluidSimulation.m_temperatureAmount);
                clearShader.SetFloat("densityAmount", fluidSimulation.m_densityAmount);
                clearShader.SetFloat("kinectAmount", kinectAmount);
                clearShader.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);
            Graphics.DrawProcedural(MeshTopology.Points, PositionBuffer.count);
        }

        private void OnDestroy()
        {
            foreach (var buf in kinectBuffers)
                buf.Release();
            clear_grid_buffer_args_buffer.Release();
            args_buffer.Release();
            KinectGridBuffer.Release();
        }

    }

}