【Unity】float型配列をカラーマップに変換するComputeShader

とある数値データを可視化する際、以下のようなカラーマップで表現することがあります。最近は街中のモニター体温計などでよく見ますね。

f:id:a_hancho:20201227011822p:plain — サーモグラフィ

f:id:a_hancho:20201227011949p:plain — 海水温

Unityでこれをやりたい場面があったので、数値を画像に変えるコンピュートシェーダーをかいてみたのでメモです。

実装

ComputeShader

#pragma kernel GetColorMapTex

Texture2D<float4> _ColorMapTex; // カラーマップの色の参考テクスチャ
int _ColorMapTexWidth;
float _CutOffValue; // ある数値以下の値は透明にしたい場合に使用
int _ResolutionX; // ColorMapの画質(横)
int _ResolutionY; // ColorMapの画質(縦)
RWTexture2D<float4> _ResultTex;

StructuredBuffer<float> _ValueBuffer;

[numthreads(1, 1 ,1)]
void GetColorMapTex(uint2 id : SV_DispatchThreadID)
{
    int valueBufferId = id.x + id.y * _ResolutionX;
    float powerVal = _ValueBuffer[valueBufferId];

    uint2 uv = uint2(
        floor(_ColorMapTexWidth * powerVal),
        0
    );
    float4 powerColor = _ColorMapTex[uv];
    if (powerVal < _CutOffValue) powerColor.a = 0;

    _ResultTex[id] = powerColor;
}

CPU側(C#)

using System;
using Cysharp.Threading.Tasks;
using System.Runtime.InteropServices;
using UnityEngine;

[RequireComponent(typeof(MeshRenderer))]
public class PowerMapVisualizer : MonoBehaviour
{
    [SerializeField] private ComputeShader _shaderSource;
    [SerializeField] private Texture _colorMapTex;
    [SerializeField] private string _texPropertyName = "_MainTex";
    [SerializeField] private float _cutOffValue = 0.1f;
    [SerializeField] private int _resolutionX = 33;
    [SerializeField] private int _resolutionY = 33;
    [SerializeField] private float _intervalSec = 0.1f;

    private float[] _powerValueArr;
    private int _shaderKernelId;
    private ComputeShader _shader;
    private ComputeBuffer _shaderBuffer;
    private RenderTexture _resultRenderTex;
    private Material _material;

    private void Start()
    {
        int powerValueLength = _resolutionX * _resolutionY;
        _powerValueArr = new float[powerValueLength];

        _shader = Instantiate(_shaderSource);
        _shaderKernelId = _shader.FindKernel("GetPowerMapTex");
        _shaderBuffer = new ComputeBuffer(powerValueLength, Marshal.SizeOf(typeof(float)));

        _resultRenderTex =
            new RenderTexture(_resolutionX, _resolutionY, 0, RenderTextureFormat.ARGB32);
        _resultRenderTex.enableRandomWrite = true;
        _resultRenderTex.Create();

        _shader.SetTexture(_shaderKernelId, "_ColorMapTex", _colorMapTex);
        _shader.SetInt("_ColorMapTexWidth", _colorMapTex.width);
        _shader.SetFloat("_CutOffValue", _cutOffValue);
        _shader.SetInt("_ResolutionX", _resolutionX);
        _shader.SetInt("ResolutionY", _resolutionY);
        _shader.SetTexture(_shaderKernelId, "_ResultRWTex", _resultRenderTex);

        _material = GetComponent<MeshRenderer>().material;
        UpdatePowerMapLoop();
    }

    private async void UpdatePowerMapLoop()
    {
        while (true)
        {
            UpdateTexture();
            await UniTask.Delay(TimeSpan.FromSeconds(_intervalSec));
        }
    }

    private void UpdateTexture()
    {
        // 仮数値データ生成箇所
        for (int y = 0; y < _resolutionY; y++)
        {
            for (int x = 0; x < _resolutionX; x++)
            {
                int index = x + y * _resolutionX;
                float lengthX = x * 2f / _resolutionX - 1f;
                float lengthY = y * 2f / _resolutionY - 1f;
                Vector2 distFromCenter = new Vector2(lengthX, lengthY);
                float length = distFromCenter.magnitude;
                _powerValueArr[index] = Mathf.Abs(Mathf.Sin(length - Time.time));
            }
        }

        _shaderBuffer.SetData(_powerValueArr);

        _shader.SetBuffer(_shaderKernelId, "_ValueBuffer", _shaderBuffer);
        _shader.Dispatch(_shaderKernelId, _resolutionX, _resolutionY, 1);
        _material.SetTexture(_texPropertyName, _resultRenderTex);
    }

    private void OnDisable()
    {
        _shaderBuffer.Release();
    }
}