anki-3d-engine/AnKi/Shaders/DepthAwareBlur.ankiprog

// Copyright (C) 2009-present, Panagiotis Christopoulos Charitos and contributors.
// All rights reserved.
// Code licensed under the BSD License.
// http://www.anki3d.org/LICENSE

#pragma anki mutator ORIENTATION 0 1 2 // 0: VERTICAL, 1: HORIZONTAL, 2: BOX
#pragma anki mutator SAMPLE_COUNT 3 5 7 9 11 13 15
#pragma anki mutator COLOR_COMPONENTS 4 3 1

#pragma anki technique vert pixel comp

#include <AnKi/Shaders/QuadVert.hlsl>
#include <AnKi/Shaders/Common.hlsl>
#include <AnKi/Shaders/BilateralFilter.hlsl>

#if ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER
#	define ORIENTATION_VERTICAL 0
#	define ORIENTATION_HORIZONTAL 1
#	define ORIENTATION_BOX 2

#	if SAMPLE_COUNT < 3
#		error See file
#	endif

// Define some macros depending on the number of components
#	if COLOR_COMPONENTS == 4
typedef Vec4 ColorType;
#	elif COLOR_COMPONENTS == 3
typedef Vec3 ColorType;
#	elif COLOR_COMPONENTS == 1
typedef F32 ColorType;
#	else
#		error See file
#	endif

SamplerState g_linearAnyClampSampler : register(s0);
Texture2D<ColorType> g_inTex : register(t0);
Texture2D<Vec4> g_depthTex : register(t1);

#	if ANKI_COMPUTE_SHADER
#		define THREADGROUP_SQRT_SIZE 8
RWTexture2D<ColorType> g_outImg : register(u0);
#	endif

F32 computeDepthWeight(F32 refDepth, F32 depth)
{
	const F32 diff = abs(refDepth - depth);
	const F32 weight = 1.0 / (kEpsilonF32 + diff);
	return sqrt(weight);
}

F32 readDepth(Vec2 uv)
{
	return g_depthTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0).r;
}

void sampleTex(Vec2 uv, F32 refDepth, inout ColorType col, inout F32 weight)
{
	const ColorType color = g_inTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0);
	const F32 w = calculateBilateralWeightDepth(refDepth, readDepth(uv), 1.0f);
	col += color * w;
	weight += w;
}

#	if ANKI_COMPUTE_SHADER
[numthreads(THREADGROUP_SQRT_SIZE, THREADGROUP_SQRT_SIZE, 1)] void main(UVec2 svDispatchThreadId : SV_DISPATCHTHREADID)
#	else
ColorType main(VertOut input) : SV_TARGET0
#	endif
{
	UVec2 textureSize;
	U32 mipCount;
	g_inTex.GetDimensions(0, textureSize.x, textureSize.y, mipCount);

	// Set UVs
#	if ANKI_COMPUTE_SHADER
	const Vec2 uv = (Vec2(svDispatchThreadId) + 0.5) / Vec2(textureSize);
#	else
	const Vec2 uv = input.m_uv;
#	endif

	const Vec2 texelSize = 1.0 / Vec2(textureSize);

	// Sample
	ColorType color = g_inTex.SampleLevel(g_linearAnyClampSampler, uv, 0.0);
	const F32 refDepth = readDepth(uv);
	F32 weight = 1.0;

#	if ORIENTATION != ORIENTATION_BOX
	// Do seperable

#		if ORIENTATION == ORIENTATION_HORIZONTAL
#			define X_OR_Y x
#		else
#			define X_OR_Y y
#		endif

	Vec2 uvOffset = 0.0f;
	uvOffset.X_OR_Y = 1.0f * texelSize.X_OR_Y;

	[unroll] for(U32 i = 0u; i < (SAMPLE_COUNT - 1u) / 2u; ++i)
	{
		sampleTex(uv + uvOffset, refDepth, color, weight);
		sampleTex(uv - uvOffset, refDepth, color, weight);

		uvOffset.X_OR_Y += 1.0f * texelSize.X_OR_Y;
	}
#	else
	// Do box

	const Vec2 offset = 1.5 * texelSize;

	sampleTex(uv + Vec2(+offset.x, +offset.y), refDepth, color, weight);
	sampleTex(uv + Vec2(+offset.x, -offset.y), refDepth, color, weight);
	sampleTex(uv + Vec2(-offset.x, +offset.y), refDepth, color, weight);
	sampleTex(uv + Vec2(-offset.x, -offset.y), refDepth, color, weight);

	sampleTex(uv + Vec2(offset.x, 0.0), refDepth, color, weight);
	sampleTex(uv + Vec2(0.0, offset.y), refDepth, color, weight);
	sampleTex(uv + Vec2(-offset.x, 0.0), refDepth, color, weight);
	sampleTex(uv + Vec2(0.0, -offset.y), refDepth, color, weight);
#	endif

	color /= weight;

	// Write value
#	if ANKI_COMPUTE_SHADER
	g_outImg[svDispatchThreadId] = color;
#	else
	return color;
#	endif
}
#endif // ANKI_COMPUTE_SHADER || ANKI_PIXEL_SHADER