環境

Unity2021.2.18f1

概要

DeepNudeの学習済みデータが手に入ったので、onnxファイルに変換してBarracudaで動作確認してみました。

せっかくなので、その情報（水着部分を緑色で塗る）を使い、ComputeShaderで水玉コラージュを作成してみました。

顔情報の取得にはUltraFaceBarracudaを使用させていただいています。

こちらのサイトの無料素材を使用させていただいています。 www.photo-ac.com

入力画像としては512x512ではないとダメなのでそうなるように画像ソフトであらかじめ修正をしています。

float4 FaceRect;
uint Width;
uint Height;
uint GridSize;
uint GridWidth;
uint GridHeight;
float4 MaskColor;
RWTexture2D<float4> Result;
Texture2D<float4> SrcTex;
RWStructuredBuffer<uint> Grids;
RWStructuredBuffer<uint2> Masks;
RWStructuredBuffer<uint> MaskCount;
RWStructuredBuffer<float> Distances;
float4 Pos;

#pragma kernel CSWriteFaceRect
[numthreads(8,8,1)]
void CSWriteFaceRect(uint3 id : SV_DispatchThreadID)
{
    if((uint)FaceRect.x <= id.x && id.x <= (uint)FaceRect.z && (uint)FaceRect.y <= id.y && id.y <= (uint)FaceRect.w)
        Result[id.xy] = float4(1, 1, 1, 1);
}

#pragma kernel CSGrid
[numthreads(8, 8, 1)]
void CSGrid(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= GridWidth)
        return;
    if (id.y >= GridHeight)
        return;
    uint2 basePos = id.xy * GridSize;
    uint value = 0;
    for (uint y = 0; y < GridSize; y++)
    {
        uint py = basePos.y + y;
        if (py >= Height)
            break;
        for (uint x = 0; x < GridSize; x++)
        {
            uint2 pos = uint2(basePos.x + x, py);
            if (pos.x >= Width)
                break;
            float4 col = SrcTex[pos.xy];
            if (col.r == MaskColor.r && col.g == MaskColor.g && col.b == MaskColor.b)
            {
                value = 1;
                break;
            }
        }
    }
    Grids[id.y * GridWidth + id.x] = (value > 0) ? 1 : 0;
}

#pragma kernel CSWriteGrid
[numthreads(8, 8, 1)]
void CSWriteGrid(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= Width)
        return;
    if (id.y >= Height)
        return;
    uint2 gridPos = id.xy / GridSize;
    uint mask = Grids[gridPos.y * GridWidth + gridPos.x];
    Result[id.xy] = (mask == 1) ? float4(0, 0, 1, 1) : Result[id.xy];
}

#pragma kernel CSClear
[numthreads(512, 1, 1)]
void CSClear(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= Width)
        return;
    if (id.y >= Height)
        return;
    Result[id.xy] = float4(0, 0, 0, 1);
}

#pragma kernel CSCopy
[numthreads(512, 1, 1)]
void CSCopy(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= Width)
        return;
    if (id.y >= Height)
        return;
    Result[id.xy] = SrcTex[id.xy];
}

#pragma kernel CSCombine
[numthreads(512, 1, 1)]
void CSCombine(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= Width)
        return;
    if (id.y >= Height)
        return;
    Result[id.xy] = lerp(float4(1, 0, 1, 1), SrcTex[id.xy], Result[id.xy].r);
}

#pragma kernel CSGridMask
[numthreads(1, 1, 1)]
void CSGridMask(uint3 id : SV_DispatchThreadID)
{
    MaskCount[0] = 0;
    uint count = GridWidth * GridHeight;
    for (uint i = 0; i < count; i++)
    {
        if (Grids[i] != 1)
            continue;
        uint x = i % GridWidth;
        uint y = i / GridWidth;
        Masks[MaskCount[0]] = uint2(x, y);
        MaskCount[0] += 1;
    }
}

#pragma kernel CSInitGridDistance
[numthreads(512, 1, 1)]
void CSInitGridDistance(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= GridWidth)
        return;
    if (id.y >= GridHeight)
        return;
    uint index = id.y * GridWidth + id.x;
    uint count = GridWidth * GridHeight;
    Distances[index] = float(count);
}

#pragma kernel CSGridDistance
[numthreads(8, 8, 1)]
void CSGridDistance(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= GridWidth)
        return;
    if (id.y >= GridHeight)
        return;
    uint index = id.y * GridWidth + id.x;
    for (uint i = 0; i < MaskCount[0]; i++)
    {
        int2 v = Masks[i] - id.xy;
        float distance = max(length(v) - 1, 0);
        if (Distances[index] > distance)
            Distances[index] = distance;
    }
}

#pragma kernel CSWriteGridDistance
[numthreads(8, 8, 1)]
void CSWriteGridDistance(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= Width)
        return;
    if (id.y >= Height)
        return;
    uint2 gridPos = id.xy / GridSize;
    uint index = gridPos.y * GridWidth + gridPos.x;
    float c = (Distances[index] == 0) ? 0 : 1;
    Result[id.xy] = float4(c, c, c, 1);
}

#pragma kernel CSWriteCircle
[numthreads(8, 8, 1)]
void CSWriteCircle(uint3 id : SV_DispatchThreadID)
{
    if (id.x >= Width)
        return;
    if (id.y >= Height)
        return;
    uint2 gridPos = int2(Pos.xy);
    int2 pos = gridPos.xy * GridSize;
    float distance = length(int2(id.xy) - pos);
    uint index = gridPos.y * GridWidth + gridPos.x;
    if (distance >= (Distances[index] * GridSize))
        return;
    Result[id.xy] = float4(1, 1, 1, 1);

    gridPos = id.xy / GridSize;
    index = gridPos.y * GridWidth + gridPos.x;
    Grids[index] = 1;
}

using UnityEngine;
using Unity.Barracuda;
using UltraFace;
using System.Linq;
using System;
using UnityEngine.UI;
using System.Collections.Generic;

public unsafe class GenerateMizutamaCollage : MonoBehaviour
{
    [SerializeField] private Texture _inputTexture = null;
    [Header("DeepNude")]
    [SerializeField] private NNModel _modelAsset = null;
    [Header("FaceDitect")]
    [SerializeField] ResourceSet _resources = null;
    [SerializeField, Range(0, 1)] float _threshold = 0.9f;
    [SerializeField] ComputeShader _computeShader = null;
    [SerializeField] RawImage _rawImageInput = null;
    [SerializeField] RawImage _rawImageOutput = null;

    private Model _runtimeModel = null;
    private IWorker _worker = null;
    private FaceDetector _detector;
    private RenderTexture _outputTexture = null;
    private RenderTexture[] _tempTextures = null;

    private static int _spResult = Shader.PropertyToID("Result");
    private static int _spFaceRect = Shader.PropertyToID("FaceRect");

    private static int _spMaskColor = Shader.PropertyToID("MaskColor");
    private static int _spWidth = Shader.PropertyToID("Width");
    private static int _spHeight = Shader.PropertyToID("Height");
    private static int _spGridSize = Shader.PropertyToID("GridSize");
    private static int _spGridWidth = Shader.PropertyToID("GridWidth");
    private static int _spGridHeight = Shader.PropertyToID("GridHeight");
    private static int _spSrcTex = Shader.PropertyToID("SrcTex");
    private static int _spGrids = Shader.PropertyToID("Grids");
    private static int _spMaskCount = Shader.PropertyToID("MaskCount");
    private static int _spMasks = Shader.PropertyToID("Masks");
    private static int _spDistances = Shader.PropertyToID("Distances");
    private static int _spPos = Shader.PropertyToID("Pos");

    private RectInt _faceRect;
    private Vector2Int _gridFaceCenter;
    private const int _gridSize = 2;
    private int _gridWidth, _gridHeight;
    private GraphicsBuffer _gbGrids = null;
    private GraphicsBuffer _gbMasks = null;
    private GraphicsBuffer _gbMaskCount = null;
    private GraphicsBuffer _gbDistances = null;

    private void Start()
    {
        _rawImageInput.texture = _inputTexture;
        _rawImageInput.rectTransform.sizeDelta = new Vector2(_inputTexture.width * 0.25f, _inputTexture.height * 0.25f);

        _outputTexture = new RenderTexture(_inputTexture.width, _inputTexture.height, 0, RenderTextureFormat.ARGB32, 0);
        _outputTexture.enableRandomWrite = true;
        _rawImageOutput.texture = _outputTexture;
        _rawImageOutput.rectTransform.sizeDelta = new Vector2(_outputTexture.width, _outputTexture.height);

        _tempTextures = new RenderTexture[2];
        for(int i = 0; i < _tempTextures.Length; i++)
        {
            _tempTextures[i] = new RenderTexture(_inputTexture.width, _inputTexture.height, 0, RenderTextureFormat.ARGB32, 0);
            _tempTextures[i].enableRandomWrite = true;
        }

        _runtimeModel = ModelLoader.Load(_modelAsset);
        _worker = WorkerFactory.CreateWorker(WorkerFactory.Type.Compute, _runtimeModel);
        Execute();

#if True
        //FaceDetect
        _detector = new FaceDetector(_resources);
        _detector.ProcessImage(_inputTexture, _threshold);
        foreach (var d in _detector.Detections)
            Debug.Log(d);
        var ditection = _detector.Detections.FirstOrDefault();
        _gridFaceCenter = Vector2Int.zero;
        var isDetect = false;
        if(ditection.score > _threshold)
        {
            isDetect = true;
            _faceRect = new RectInt(
                (int)(ditection.x1 * _inputTexture.width),
                (int)((1.0f - ditection.y2) * _inputTexture.height),
                (int)((ditection.x2 - ditection.x1) * _inputTexture.width),
                (int)((ditection.y2 - ditection.y1) * _inputTexture.height)
            );
            Debug.Log($"faceRect:{_faceRect}");
            _gridFaceCenter = new Vector2Int((int)(_faceRect.center.x / (float)_gridSize), (int)(_faceRect.center.y / (float)_gridSize));
        }
        CopyTexture(_outputTexture, _tempTextures[0]);
        //DrawFace();
        CreateGrids();
        //DrawGrids();
        ClearTexture();
        _gbMasks = new GraphicsBuffer(GraphicsBuffer.Target.Structured, _gridWidth * _gridHeight, sizeof(Vector2Int));
        _gbMaskCount = new GraphicsBuffer(GraphicsBuffer.Target.Structured, 1, sizeof(UInt32));
        _gbDistances = new GraphicsBuffer(GraphicsBuffer.Target.Structured, _gridWidth * _gridHeight, sizeof(float));
        var corners = new Vector2Int[]
        {
            new Vector2Int(_gridFaceCenter.x, _gridFaceCenter.y),
            new Vector2Int(0, 0),
            new Vector2Int(_gridWidth - 1, 0),
            new Vector2Int(_gridWidth - 1, _gridHeight - 1),
            new Vector2Int(0, _gridHeight - 1),
        };
        for(int i = (isDetect == true) ? 0 : 1; i < corners.Length; i++)
        {
            CreateMaskPositions();
            CreateDistances();
            //DrawDistances(_tempTextures[1]);
            DrawCircle(corners[i]);
        }
        CreateMaskPositions();
        CreateDistances();

        var grids = new UInt32[_gridWidth * _gridHeight];
        _gbGrids.GetData(grids);
        var positions = new List<Vector2Int>();
        for(int i = 0; i < _gridWidth * _gridHeight; i++)
        {
            var x = i % _gridWidth;
            var y = i / _gridWidth;
            if(grids[i] == 1)
                continue;
            positions.Add(new Vector2Int(x, y));
        }
        //シャッフル
        positions = positions.OrderBy(a => Guid.NewGuid()).ToList();
        var count = Mathf.Min(200, positions.Count);
        for(int i = 0; i < count; i++)
        {
            var pos = positions[i];
            DrawCircle(pos);
            CreateMaskPositions();
            CreateDistances();
        }
        CombineTexture(_inputTexture);

        _gbMasks.Dispose();
        _gbMaskCount.Dispose();
        _gbDistances.Dispose();
        _gbGrids.Dispose();
#endif
    }

    private void Execute()
    {
        Tensor input = new Tensor(_inputTexture as Texture, channels:3);
        Inference(input);
        input.Dispose();
    }

    private void Inference(Tensor input)
    {
        _worker.Execute(input);
        Tensor output = _worker.PeekOutput();
        output.ToRenderTexture(_outputTexture, 0, 0, 1, 0, null);
        output.Dispose();
    }

    private void DrawFace()
    {
        //顔Rect描画
        var kernelIndex = _computeShader.FindKernel("CSWriteFaceRect");
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.SetTexture(kernelIndex, _spResult, _outputTexture);
        _computeShader.SetVector(_spFaceRect, new Vector4(_faceRect.xMin, _faceRect.yMin, _faceRect.xMax, _faceRect.yMax));
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_outputTexture.width / (float)numThreadX), Mathf.CeilToInt((float)_outputTexture.height / (float)numThreadY), 1);
    }
    private void CreateGrids()
    {
        //グリッドで分割
        _gridWidth = Mathf.CeilToInt((float)_outputTexture.width / (float)_gridSize);
        _gridHeight = Mathf.CeilToInt((float)_outputTexture.height / (float)_gridSize);
        _gbGrids = new GraphicsBuffer(GraphicsBuffer.Target.Structured, _gridWidth * _gridHeight, sizeof(UInt32));
        var kernelIndex = _computeShader.FindKernel("CSGrid");
        _computeShader.SetVector(_spMaskColor, Color.green);
        _computeShader.SetInt(_spWidth, _outputTexture.width);
        _computeShader.SetInt(_spHeight, _outputTexture.height);
        _computeShader.SetInt(_spGridSize, _gridSize);
        _computeShader.SetInt(_spGridWidth, _gridWidth);
        _computeShader.SetInt(_spGridHeight, _gridHeight);
        _computeShader.SetTexture(kernelIndex, _spSrcTex, _outputTexture);
        _computeShader.SetBuffer(kernelIndex, _spGrids, _gbGrids);
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_gridWidth / (float)numThreadX), Mathf.CeilToInt((float)_gridHeight / (float)numThreadY), 1);
    }
    private void DrawGrids()
    {
        //グリッド描画
        var kernelIndex = _computeShader.FindKernel("CSWriteGrid");
        _computeShader.SetBuffer(kernelIndex, _spGrids, _gbGrids);
        _computeShader.SetTexture(kernelIndex, _spResult, _outputTexture);
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_outputTexture.width / (float)numThreadX), Mathf.CeilToInt((float)_outputTexture.height / (float)numThreadY), 1);
    }
    private void ClearTexture()
    {
        var kernelIndex = _computeShader.FindKernel("CSClear");
        _computeShader.SetInt(_spWidth, _outputTexture.width);
        _computeShader.SetInt(_spHeight, _outputTexture.height);
        _computeShader.SetTexture(kernelIndex, _spResult, _outputTexture);
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_outputTexture.width / (float)numThreadX), Mathf.CeilToInt((float)_outputTexture.height / (float)numThreadY), 1);
    }
    private void CopyTexture(Texture src, Texture dst)
    {
        var kernelIndex = _computeShader.FindKernel("CSCopy");
        _computeShader.SetInt(_spWidth, _outputTexture.width);
        _computeShader.SetInt(_spHeight, _outputTexture.height);
        _computeShader.SetTexture(kernelIndex, _spSrcTex, src);
        _computeShader.SetTexture(kernelIndex, _spResult, dst);
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_outputTexture.width / (float)numThreadX), Mathf.CeilToInt((float)_outputTexture.height / (float)numThreadY), 1);
    }
    private void CombineTexture(Texture src)
    {
        var kernelIndex = _computeShader.FindKernel("CSCombine");
        _computeShader.SetInt(_spWidth, _outputTexture.width);
        _computeShader.SetInt(_spHeight, _outputTexture.height);
        _computeShader.SetTexture(kernelIndex, _spSrcTex, src);
        _computeShader.SetTexture(kernelIndex, _spResult, _outputTexture);
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_outputTexture.width / (float)numThreadX), Mathf.CeilToInt((float)_outputTexture.height / (float)numThreadY), 1);
    }
    private void CreateMaskPositions()
    {
        var kernelIndex = _computeShader.FindKernel("CSGridMask");
        _computeShader.SetInt(_spGridWidth, _gridWidth);
        _computeShader.SetInt(_spGridHeight, _gridHeight);
        _computeShader.SetBuffer(kernelIndex, _spGrids, _gbGrids);
        _computeShader.SetBuffer(kernelIndex, _spMasks, _gbMasks);
        _computeShader.SetBuffer(kernelIndex, _spMaskCount, _gbMaskCount);
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.Dispatch(kernelIndex, 1,1,1);
    }
    private void CreateDistances()
    {
        //距離マップの作成
        var kernelIndex = _computeShader.FindKernel("CSInitGridDistance");
        _computeShader.SetInt(_spGridWidth, _gridWidth);
        _computeShader.SetInt(_spGridHeight, _gridHeight);
        _computeShader.SetBuffer(kernelIndex, _spDistances, _gbDistances);
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_gridWidth / (float)numThreadX), Mathf.CeilToInt((float)_gridHeight / (float)numThreadY), 1);

        kernelIndex = _computeShader.FindKernel("CSGridDistance");
        _computeShader.SetInt(_spGridWidth, _gridWidth);
        _computeShader.SetInt(_spGridHeight, _gridHeight);
        _computeShader.SetInt(_spGridSize, _gridSize);
        _computeShader.SetBuffer(kernelIndex, _spMasks, _gbMasks);
        _computeShader.SetBuffer(kernelIndex, _spMaskCount, _gbMaskCount);
        _computeShader.SetBuffer(kernelIndex, _spDistances, _gbDistances);
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_gridWidth / (float)numThreadX), Mathf.CeilToInt((float)_gridHeight / (float)numThreadY), 1);
    }
    private void DrawDistances(Texture dst)
    {
        //距離マップ描画
        var kernelIndex = _computeShader.FindKernel("CSWriteGridDistance");
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.SetInt(_spWidth, _outputTexture.width);
        _computeShader.SetInt(_spHeight, _outputTexture.height);
        _computeShader.SetInt(_spGridSize, _gridSize);
        _computeShader.SetInt(_spGridWidth, _gridWidth);
        _computeShader.SetTexture(kernelIndex, _spResult, dst);
        _computeShader.SetBuffer(kernelIndex, _spDistances, _gbDistances);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_outputTexture.width / (float)numThreadX), Mathf.CeilToInt((float)_outputTexture.height / (float)numThreadY), 1);
    }
    private void DrawCircle(Vector2Int pos)
    {
        //円の描画
        var kernelIndex = _computeShader.FindKernel("CSWriteCircle");
        uint numThreadX, numThreadY, numThreadZ;
        _computeShader.GetKernelThreadGroupSizes(kernelIndex, out numThreadX, out numThreadY, out numThreadZ);
        _computeShader.SetInt(_spWidth, _outputTexture.width);
        _computeShader.SetInt(_spHeight, _outputTexture.height);
        _computeShader.SetInt(_spGridSize, _gridSize);
        _computeShader.SetInt(_spGridWidth, _gridWidth);
        _computeShader.SetVector(_spPos, new Vector4(pos.x, pos.y, 0,0));
        _computeShader.SetTexture(kernelIndex, _spResult, _outputTexture);
        _computeShader.SetBuffer(kernelIndex, _spDistances, _gbDistances);
        _computeShader.SetBuffer(kernelIndex, _spGrids, _gbGrids);
        _computeShader.Dispatch(kernelIndex, Mathf.CeilToInt((float)_outputTexture.width / (float)numThreadX), Mathf.CeilToInt((float)_outputTexture.height / (float)numThreadY), 1);
    }

    private void OnDestroy()
    {
        _detector?.Dispose();
        _worker?.Dispose();
        _outputTexture.Release();
        if(_tempTextures != null)
        {
            for(int i = 0; i < _tempTextures.Length; i++)
                _tempTextures[i].Release();
        }
    }
}