CLAHE - Contrast Limited Adaptive Histogram Normalization

/** Local Contrast PS, version 0.1.0
All rights (c) 2020 Jakub Maksymilian Fober (the Author).

The Author provides this shader (the Work)
under the Creative Commons CC BY-SA 3.0 license
available online at
http://creativecommons.org/licenses/by-sa/3.0/

For inquiries please contact jakub.m.fober@pm.me
*/


#include "ReShadeUI.fxh"
#include "ReShade.fxh"


  ////////////
 /// MENU ///
////////////

#if !defined(BLOCK_SIZE)
	#define BLOCK_SIZE 8
#endif

uniform bool CorrectGamma <
	ui_label = "Gamma correct";
	ui_tootlip = "Applies S-curve to the corrected luma channel";
> = true;

uniform float ContrastLimit < __UNIFORM_SLIDER_FLOAT1
	ui_label = "Contrast limit";
	ui_min = 1.0; ui_max = 0.5;
> = 0.75;


  ///////////////
 /// TEXTURE ///
///////////////

// Local histogram values map
texture LocalContrastMapBuffer
{
	Width = BUFFER_WIDTH/BLOCK_SIZE;
	Height = BUFFER_HEIGHT/BLOCK_SIZE;
	Format = RG8;
	// R = histogram luma min
	// G = histogram luma max
};
sampler LocalContrastMap { Texture = LocalContrastMapBuffer; };

// Point mapping back-buffer sampler
sampler BackBuffer { Texture = ReShade::BackBufferTex; };

  /////////////////
 /// FUNCTIONS ///
/////////////////

// RGB to YUV709
static const float3x3 ToYUV709 =
float3x3(
	float3(0.2126, 0.7152, 0.0722),
	float3(-0.09991, -0.33609, 0.436),
	float3(0.615, -0.55861, -0.05639)
);
// RGB to YUV601
static const float3x3 ToYUV601 =
float3x3(
	float3(0.299, 0.587, 0.114),
	float3(-0.14713, -0.28886, 0.436),
	float3(0.615, -0.51499, -0.10001)
);
// YUV709 to RGB
static const float3x3 ToRGB709 =
float3x3(
	float3(1, 0, 1.28033),
	float3(1, -0.21482, -0.38059),
	float3(1, 2.12798, 0)
);
// YUV601 to RGB
static const float3x3 ToRGB601 =
float3x3(
	float3(1, 0, 1.13983),
	float3(1, -0.39465, -0.58060),
	float3(1, 2.03211, 0)
);

// Convert linear weight to gaussian
float weight(float x)
{
	const float PI = radians(180);
	x -= 0.5;
	return sin(x*PI)*0.5+0.5;
}
float2 mapCoord(float2 offset, float2 texcoord)
{ return texcoord*(1.0-offset)+offset*0.5; }

  //////////////
 /// SHADER ///
//////////////

// Analyze histogram per image block
void GetLocalHistogramPS(
	float4 pos : SV_Position,
	float2 texcoord : TEXCOORD,
	out float2 histogramStats : SV_Target // Contrast-limited adaptive histogram
){
	// Convert texture coordinates bounds
	// texcoord = mapCoord(
	// 	BUFFER_PIXEL_SIZE,
	// 	mapCoord(rcp(BUFFER_SCREEN_SIZE/8), texcoord)
	// )+BUFFER_PIXEL_SIZE*0.5;
	// Initial values of the histogram
	histogramStats.s = 1.0; // Min
	histogramStats.t = 0.0; // Max
	float histogramMean = 0.0; // Average
	// Loop through image block and get histogram min, max and average
	const int halfBlock = BLOCK_SIZE/2;
	for (int y=-halfBlock; y<halfBlock; y++)
	for (int x=-halfBlock; x<halfBlock; x++)
	{
		// Get block background picture texture coordinates
		float2 blockCoord = BUFFER_PIXEL_SIZE*float2(x, y)+texcoord;
		// Get luminosity of background picture pixel of a block
		float luma = dot(ToYUV709[0], tex2D(BackBuffer, blockCoord).rgb);
		// Save histogram data of the block
		histogramStats.s = min(luma, histogramStats.s); // Min
		histogramStats.t = max(luma, histogramStats.t); // Max
		histogramMean += luma; // Average
	}
	histogramMean /= BLOCK_SIZE*BLOCK_SIZE;
	// Contrast limiting
	histogramStats.s = min(histogramStats.s, max(0.0, histogramMean-ContrastLimit));
	histogramStats.t = max(histogramStats.t, min(1.0, histogramMean+ContrastLimit));
}


// Apply local contrast to image
void LocalConstrastPS(
	float4 pos : SV_Position,
	float2 texcoord : TEXCOORD,
	out float3 result : SV_Target
){
	// Get background color in YUV
	result = mul(ToYUV709, tex2D(BackBuffer, texcoord).rgb);
	// Get local contrast map
	float2 localContrast = tex2D(LocalContrastMap, texcoord).st;
	// Histogram normalization of luma channel
	result.s = (result.s-localContrast.s)/(localContrast.t-localContrast.s);
	// S-curve correction
	if (CorrectGamma)
	{
		float s_curve = 2.0-ContrastLimit*2.0;
		s_curve *= s_curve*s_curve; // Power 3
		result.s = lerp(result.s, weight(result.s), s_curve);
	}
	// Convert to RGB
	result = mul(ToRGB709, result);
}


  //////////////
 /// OUTPUT ///
//////////////

technique LocalContrast <
	ui_label = "Local Contrast";
	ui_tooltip =
		"CLAHE (contrast-limited adaptive histogram normalization)\n"
		"\n"
		"To change block size, edit global preprocessor definition:\n"
		"\tBLOCK_SIZE\tdefault is 8"; >
{
	pass Local_Histogram
	{
		VertexShader = PostProcessVS;
		PixelShader = GetLocalHistogramPS;
		RenderTarget = LocalContrastMapBuffer;
	}
	pass Historgram_Normalization
	{
		VertexShader = PostProcessVS;
		PixelShader = LocalConstrastPS;
	}
}