【Unity】ComputeShaderで平均プーリングをする【Q7】

はじめに

画像処理100本ノックの続きをやってきましょう。

github.com

前回 -> 【Unity】ComputeShaderで減色処理(ポスタライズ)をしてみる【Q6】 - はなちるのマイノート

次回 -> 【Unity】ComputeShaderでMaxプーリングをする【Q8】 - はなちるのマイノート

今回は平均プーリングについて取り組んでいきます。

f:id:hanaaaaaachiru:20200208221652p:plain

はじめに
平均プーリングとは
コード
さいごに

平均プーリングとは

平均プーリングとは8*8ピクセルといったグリッド単位で画像を分割してその領域内を平均値で埋めることです。

これをGPUを用いてするので、今回は次の画像のような方針でやってみたいと思います。

f:id:hanaaaaaachiru:20200208214255p:plain

実際の例で説明をすると、入力画像が512*512ピクセルとし、ひとつのグリッドを8*8ピクセルだとすると64(=512/8)*64個のスレッドを用いて計算をします。

f:id:hanaaaaaachiru:20200208215140p:plain

コード

computeShaderはこちら。

#pragma kernel AveragePooling

RWTexture2D<float4> Result;
Texture2D<float4> Texture;
int2 Grid;

[numthreads(8,8,1)]
void AveragePooling (uint3 id : SV_DispatchThreadID)
{
    float3 sum = float3(0,0,0);

    for(int i = 0; i < Grid.x; i++){
        for(int j = 0; j < Grid.y; j++){
            int2 index = int2((id.x * Grid.x) + i, (id.y * Grid.y) + j);
            sum.x += Texture[index.xy].x;
            sum.y += Texture[index.xy].y;
            sum.z += Texture[index.xy].z;
        }
    }

    float r = sum.x / (Grid.x * Grid.y);;
    float g = sum.y / (Grid.x * Grid.y);
    float b = sum.z / (Grid.x * Grid.y);
    
    for(int i = 0; i < Grid.x; i++){
        for(int j = 0; j < Grid.y; j++){
            int2 index = int2((id.x * Grid.x) + i, (id.y * Grid.y) + j);
            Result[index.xy] = float4(r, g, b, 1);
        }
    }
}

CPU側のコードはこちら。

using UnityEngine;
using UnityEngine.UI;

public class AveragePooling : MonoBehaviour
{
    [SerializeField] private ComputeShader _computeShader;
    [SerializeField] private Texture2D _tex;
    [SerializeField] private RawImage _renderer;

    struct ThreadSize
    {
        public uint x;
        public uint y;
        public uint z;

        public ThreadSize(uint x, uint y, uint z)
        {
            this.x = x;
            this.y = y;
            this.z = z;
        }
    }

    private void Start()
    {
        if (!SystemInfo.supportsComputeShaders)
        {
            Debug.LogError("Comppute Shader is not support.");
            return;
        }

        // グリッドを「8ピクセル＊8ピクセル」で分割
        int[] grid = new int[] { 8, 8 };

        // RenderTextueの初期化
        var result = new RenderTexture(_tex.width, _tex.height, 0, RenderTextureFormat.ARGB32);
        result.enableRandomWrite = true;
        result.Create();

        // AveragePoolingのカーネルインデックス(0)を取得
        var kernelIndex = _computeShader.FindKernel("AveragePooling");

        // 一つのグループの中に何個のスレッドがあるか
        ThreadSize threadSize = new ThreadSize();
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out threadSize.x, out threadSize.y, out threadSize.z);

        // GPUにデータをコピーする
        _computeShader.SetTexture(kernelIndex, "Texture", _tex);
        _computeShader.SetTexture(kernelIndex, "Result", result);
        _computeShader.SetInts("Grid", grid);

        // GPUの処理を実行する
        _computeShader.Dispatch(kernelIndex, (_tex.width / (int)threadSize.x) / grid[0], (_tex.height / (int)threadSize.y) / grid[1], (int)threadSize.z);

        // テクスチャを適応
        _renderer.texture = result;
    }
}