Merge pull request #421 from w23/denoiser-massage

1. Denoise different channels (diffuse, specular, indirect) differently.
2. Make indirect bounce half res.
This commit is contained in:
Ivan Avdeev 2023-02-11 12:39:07 -08:00 committed by GitHub
commit eaf8b3c55a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 159 additions and 134 deletions

View File

@ -164,9 +164,11 @@ void computeBounce(ivec2 pix, vec3 direction, out vec3 diffuse, out vec3 specula
}
}
const int INDIRECT_SCALE = 2;
void main() {
const ivec2 pix = ivec2(gl_GlobalInvocationID);
const ivec2 res = ivec2(imageSize(out_indirect_diffuse));
const ivec2 res = ivec2(imageSize(out_indirect_diffuse)) / INDIRECT_SCALE;
if (any(greaterThanEqual(pix, res))) {
return;
}
@ -179,7 +181,7 @@ void main() {
rand01_state = ubo.ubo.random_seed + pix.x * 1833 + pix.y * 31337 + 12;
vec3 diffuse, specular;
computeBounce(pix, direction, diffuse, specular);
computeBounce(pix * INDIRECT_SCALE, direction, diffuse, specular);
imageStore(out_indirect_diffuse, pix, vec4(diffuse, 0.f));
imageStore(out_indirect_specular, pix, vec4(specular, 0.f));
}

View File

@ -23,7 +23,6 @@ layout(set = 0, binding = 6, rgba16f) uniform readonly image2D emissive;
layout(set = 0, binding = 7, rgba32f) uniform readonly image2D position_t;
layout(set = 0, binding = 8, rgba16f) uniform readonly image2D normals_gs;
layout(set = 0, binding = 9, rgba16f) uniform readonly image2D prev_dest;
layout(set = 0, binding = 10, rgba32f) uniform readonly image2D geometry_prev_position;
layout(set = 0, binding = 11) uniform UBO { UniformBuffer ubo; } ubo;
@ -31,6 +30,14 @@ layout(set = 0, binding = 11) uniform UBO { UniformBuffer ubo; } ubo;
layout(set = 0, binding = 12, rgba16f) uniform readonly image2D indirect_diffuse;
layout(set = 0, binding = 13, rgba16f) uniform readonly image2D indirect_specular;
layout(set = 0, binding = 14, rgba16f) uniform image2D out_temporal_diffuse;
layout(set = 0, binding = 15, rgba16f) uniform image2D prev_temporal_diffuse;
layout(set = 0, binding = 16, rgba16f) uniform image2D out_temporal_specular;
layout(set = 0, binding = 17, rgba16f) uniform image2D prev_temporal_specular;
const int INDIRECT_SCALE = 2;
//layout(set = 0, binding = 2, rgba16f) uniform readonly image2D light_poly_diffuse;
/* layout(set = 0, binding = 3, rgba16f) uniform readonly image2D specular; */
/* layout(set = 0, binding = 4, rgba16f) uniform readonly image2D additive; */
@ -73,6 +80,102 @@ void readNormals(ivec2 uv, out vec3 geometry_normal, out vec3 shading_normal) {
shading_normal = normalDecode(n.zw);
}
struct Components {
vec3 direct_diffuse, direct_specular, indirect_diffuse, indirect_specular;
};
Components blurSamples(const ivec2 res, const ivec2 pix) {
Components c;
c.direct_diffuse = c.direct_specular = c.indirect_diffuse = c.indirect_specular = vec3(0.);
const vec4 center_pos = imageLoad(position_t, pix);
const int DIRECT_DIFFUSE_KERNEL = 3;
const int INDIRECT_DIFFUSE_KERNEL = 5;
const int SPECULAR_KERNEL = 2;
const int KERNEL_SIZE = max(max(DIRECT_DIFFUSE_KERNEL, INDIRECT_DIFFUSE_KERNEL), SPECULAR_KERNEL);
const float direct_diffuse_sigma = DIRECT_DIFFUSE_KERNEL / 2.;
const float indirect_diffuse_sigma = INDIRECT_DIFFUSE_KERNEL / 2.;
const float specular_sigma = SPECULAR_KERNEL / 2.;
float direct_diffuse_total = 0.;
float indirect_diffuse_total = 0.;
float specular_total = 0.;
const ivec2 res_scaled = res / INDIRECT_SCALE;
for (int x = -KERNEL_SIZE; x <= KERNEL_SIZE; ++x)
for (int y = -KERNEL_SIZE; y <= KERNEL_SIZE; ++y) {
const ivec2 p = pix + ivec2(x, y);
if (any(greaterThanEqual(p, res)) || any(lessThan(p, ivec2(0)))) {
continue;
}
// const vec4 c = imageLoad(light_poly, p);
// if (c.a != material_index)
// continue;
vec3 sample_geometry_normal, sample_shading_normal;
readNormals(p, sample_geometry_normal, sample_shading_normal);
float scale = 1.f;
// FIXME also filter by depth, (kusok index?), etc
//scale *= smoothstep(.9, 1., dot(sample_geometry_normal, geometry_normal));
const vec4 sample_pos = imageLoad(position_t, p);
// FIXME what are these magic numbers?
scale *= smoothstep(center_pos.w * 4. / 100., 0., distance(center_pos.xyz, sample_pos.xyz));
if ( scale <= 0. )
continue;
if (all(lessThan(abs(ivec2(x, y)), ivec2(DIRECT_DIFFUSE_KERNEL))))
{
const float direct_diffuse_scale = scale * normpdf(x, direct_diffuse_sigma) * normpdf(y, direct_diffuse_sigma);
direct_diffuse_total += direct_diffuse_scale;
c.direct_diffuse += imageLoad(light_point_diffuse, p).rgb * direct_diffuse_scale;
c.direct_diffuse += imageLoad(light_poly_diffuse, p).rgb * direct_diffuse_scale;
}
if (all(lessThan(abs(ivec2(x, y)), ivec2(INDIRECT_DIFFUSE_KERNEL))))
{
const ivec2 pscaled = pix / INDIRECT_SCALE + ivec2(x, y);
if (all(lessThan(pscaled, res_scaled)) && all(greaterThanEqual(pscaled, ivec2(0)))) {
const float indirect_diffuse_scale = scale * normpdf(x, indirect_diffuse_sigma) * normpdf(y, indirect_diffuse_sigma);
indirect_diffuse_total += indirect_diffuse_scale;
c.indirect_diffuse += imageLoad(indirect_diffuse, pscaled).rgb * indirect_diffuse_scale;
}
}
if (all(lessThan(abs(ivec2(x, y)), ivec2(SPECULAR_KERNEL))))
{
const float specular_scale = scale * normpdf(x, specular_sigma) * normpdf(y, specular_sigma);
specular_total += specular_scale;
c.direct_specular += imageLoad(light_poly_specular, p).rgb * specular_scale;
c.direct_specular += imageLoad(light_point_specular, p).rgb * specular_scale;
c.indirect_specular += imageLoad(indirect_specular, p / INDIRECT_SCALE).rgb * specular_scale;
}
}
if (direct_diffuse_total > 0.)
c.direct_diffuse /= direct_diffuse_total;
if (indirect_diffuse_total > 0.)
c.indirect_diffuse *= indirect_diffuse_total;
if (specular_total > 0.) {
c.direct_specular *= specular_total;
c.indirect_specular *= specular_total;
}
return c;
}
void main() {
ivec2 res = ivec2(imageSize(base_color_a));
ivec2 pix = ivec2(gl_GlobalInvocationID);
@ -98,6 +201,7 @@ void main() {
//imageStore(out_dest, pix, vec4((imageLoad(light_poly, pix).rgb * base_color_a.rgb), 0.)); return;
//imageStore(out_dest, pix, vec4(imageLoad(normals, pix)*.5 + .5f)); return;
//imageStore(out_dest, pix, vec4(imageLoad(specular, pix)*.5 + .5f)); return;
//imageStore(out_dest, pix, vec4(imageLoad(indirect_diffuse, pix)*.5 + .5f)); return;
//imageStore(out_dest, pix, vec4(aces_tonemap(imageLoad(light_poly, pix).rgb), 0.)); return;
//imageStore(out_dest, pix, vec4(aces_tonemap(imageLoad(specular, pix).rgb), 0.)); return;
//imageStore(out_dest, pix, vec4(aces_tonemap(imageLoad(indirect_diffuse, pix).rgb), 0.)); return;
@ -113,149 +217,68 @@ void main() {
/* imageStore(out_dest, pix, vec4(rand3_f01(uvec3(mi,mi+1,mi+2)), 0.)); */
/* return; */
#if 0
const Components c = blurSamples(res, pix);
//imageStore(out_dest, pix, vec4(aces_tonemap(c.direct_diffuse), 0.)); return;
//imageStore(out_dest, pix, vec4(aces_tonemap(c.direct_specular), 0.)); return;
//imageStore(out_dest, pix, vec4(aces_tonemap(c.indirect_diffuse), 0.)); return;
//imageStore(out_dest, pix, vec4(aces_tonemap(c.indirect_specular), 0.)); return;
vec3 colour = vec3(0.);
colour += imageLoad(light_poly_diffuse, pix).rgb;
colour += imageLoad(light_poly_specular, pix).rgb;
colour += imageLoad(light_point_diffuse, pix).rgb;
colour += imageLoad(light_point_specular, pix).rgb;
#else
float total_scale = 0.;
vec3 colour = vec3(0.);
const int KERNEL_SIZE = 2;
float specular_total_scale = 0.;
vec3 speculour = vec3(0.);
{
// DEBUG motion vectors
//colour = vec3(length(imageLoad(position_t, pix).rgb - imageLoad(prev_position_t, pix).rgb));
const vec4 center_pos = imageLoad(position_t, pix);
for (int x = -KERNEL_SIZE; x <= KERNEL_SIZE; ++x)
for (int y = -KERNEL_SIZE; y <= KERNEL_SIZE; ++y) {
const ivec2 p = pix + ivec2(x, y);
if (any(greaterThanEqual(p, res)) || any(lessThan(p, ivec2(0)))) {
continue;
}
// TODO: need to extract reprojecting from this shader because reprojected stuff need svgf denoising pass after it
const vec3 origin = (ubo.ubo.inv_view * vec4(0., 0., 0., 1.)).xyz;
const float depth = length(origin - imageLoad(position_t, pix).xyz);
const vec3 prev_position = imageLoad(geometry_prev_position, pix).rgb;
const vec4 clip_space = inverse(ubo.ubo.prev_inv_proj) * vec4((inverse(ubo.ubo.prev_inv_view) * vec4(prev_position, 1.)).xyz, 1.);
const vec2 reproj_uv = clip_space.xy / clip_space.w;
const ivec2 reproj_pix = ivec2((reproj_uv * 0.5 + vec2(0.5)) * vec2(res));
const vec3 prev_origin = (ubo.ubo.prev_inv_view * vec4(0., 0., 0., 1.)).xyz;
const float depth_nessesary = length(prev_position - prev_origin);
const float depth_treshold = 0.01 * clip_space.w;
float scale = 1.f;
float better_depth_offset = depth_treshold;
vec3 diffuse = c.direct_diffuse + c.indirect_diffuse;
vec3 specular = c.direct_specular + c.indirect_specular;
vec3 history_diffuse = diffuse;
vec3 history_specular = specular;
const int TEMPORAL_KERNEL = 1; // lifekilled says it should be fixed
for(int x = -TEMPORAL_KERNEL; x <=TEMPORAL_KERNEL; x++) {
for(int y = -TEMPORAL_KERNEL; y <=TEMPORAL_KERNEL; y++) {
const ivec2 p = reproj_pix + ivec2(x, y);
if (any(greaterThanEqual(p, res)) || any(lessThan(p, ivec2(0)))) {
continue;
}
const vec4 history_diffuse_depth = imageLoad( prev_temporal_diffuse, reproj_pix );
const vec4 history_specular_sample = imageLoad( prev_temporal_specular, reproj_pix );
// const vec4 c = imageLoad(light_poly, p);
// if (c.a != material_index)
// continue;
vec3 sample_geometry_normal, sample_shading_normal;
readNormals(p, sample_geometry_normal, sample_shading_normal);
// FIXME also filter by depth, (kusok index?), etc
//scale *= smoothstep(.9, 1., dot(sample_geometry_normal, geometry_normal));
const vec4 sample_pos = imageLoad(position_t, p);
scale *= smoothstep(4. * center_pos.w / 100., 0., distance(center_pos.xyz, sample_pos.xyz));
if ( scale <= 0. )
continue;
vec3 diffuse = vec3(0.);
diffuse += imageLoad(light_point_diffuse, p).rgb;
diffuse += imageLoad(light_poly_diffuse, p).rgb;
diffuse += imageLoad(indirect_diffuse, p).rgb;
vec3 specular = vec3(0.);
specular += imageLoad(light_poly_specular, p).rgb;
specular += imageLoad(light_point_specular, p).rgb;
specular += imageLoad(indirect_specular, p).rgb;
{
const float sigma = KERNEL_SIZE / 2.;
const float dscale = scale * normpdf(x, sigma) * normpdf(y, sigma);
colour += dscale * diffuse;
total_scale += dscale;
}
const int SPECULAR_KERNEL_SIZE = 4;
if (all(lessThan(abs(ivec2(x, y)), ivec2(SPECULAR_KERNEL_SIZE)))) {
const float spigma = SPECULAR_KERNEL_SIZE / 2.;
const float specuale = scale * normpdf(x, spigma) * normpdf(y, spigma);
speculour += specuale * specular;
specular_total_scale += specuale;
const float history_depth = history_diffuse_depth.w;
const float depth_offset = abs(history_depth - depth_nessesary);
if ( depth_offset < better_depth_offset ) {
better_depth_offset = depth_offset;
history_diffuse = history_diffuse_depth.rgb;
history_specular = history_specular_sample.rgb;
}
}
}
if (better_depth_offset < depth_treshold) {
diffuse = mix(diffuse, history_diffuse, 0.8);
specular = mix(specular, history_specular, 0.3);
}
if (total_scale > 0.) {
colour /= total_scale;
}
imageStore(out_temporal_diffuse, pix, vec4(diffuse, depth));
imageStore(out_temporal_specular, pix, vec4(specular, 0./*unused*/));
colour = diffuse + specular;
if (specular_total_scale > 0.) {
speculour /= specular_total_scale;
colour += speculour;
//imageStore(out_dest, pix, vec4(LINEARtoSRGB(diffuse), 0.)); return;
}
#endif
const vec4 base_color_a = imageLoad(base_color_a, pix);
colour *= SRGBtoLINEAR(base_color_a.rgb);
colour += imageLoad(emissive, pix).rgb;
//colour += imageLoad(additive, pix).rgb;
colour = LINEARtoSRGB(colour);
// HACK: exposure
// TODO: should be dynamic based on previous frames brightness
#if 0
if (pix.x >= res.x / 2) {
colour *= 8.;
}
#else
//colour *= .25;
#endif
//colour = aces_tonemap(colour);
//colour = reinhard02(colour, vec3(400.));
//colour = reinhard02(colour, vec3(1.));
#if 0
if (pix.x < res.x / 2) {
#endif
//colour *= .25;
colour = LINEARtoSRGB(colour); // gamma-correction
#if 0
}
#endif
// DEBUG motion vectors
//colour = vec3(length(imageLoad(position_t, pix).rgb - imageLoad(prev_position_t, pix).rgb));
// TODO: need to extract reprojecting from this shader because reprojected stuff need svgf denoising pass after it
const vec3 origin = (ubo.ubo.inv_view * vec4(0., 0., 0., 1.)).xyz;
const float depth = length(origin - imageLoad(position_t, pix).xyz);
const vec3 prev_position = imageLoad(geometry_prev_position, pix).rgb;
const vec4 clip_space = inverse(ubo.ubo.prev_inv_proj) * vec4((inverse(ubo.ubo.prev_inv_view) * vec4(prev_position, 1.)).xyz, 1.);
const vec2 reproj_uv = clip_space.xy / clip_space.w;
const ivec2 reproj_pix = ivec2((reproj_uv * 0.5 + vec2(0.5)) * vec2(res));
const vec3 prev_origin = (ubo.ubo.prev_inv_view * vec4(0., 0., 0., 1.)).xyz;
const float depth_nessesary = length(prev_position - prev_origin);
const float depth_treshold = 0.01 * clip_space.w;
float better_depth_offset = depth_treshold;
vec3 history_colour = colour;
for(int x = -1; x <=1; x++) {
for(int y = -1; y <=1; y++) {
const ivec2 p = reproj_pix + ivec2(x, y);
if (any(greaterThanEqual(p, res)) || any(lessThan(p, ivec2(0)))) {
continue;
}
const vec4 history_colour_depth = imageLoad( prev_dest, reproj_pix );
const float history_depth = history_colour_depth.w;
const float depth_offset = abs(history_depth - depth_nessesary);
if ( depth_offset < better_depth_offset ) {
better_depth_offset = depth_offset;
history_colour = history_colour_depth.rgb;
}
}
}
if (better_depth_offset < depth_treshold) {
colour = mix(colour, history_colour, 0.8);
}
// DEBUG motion vectors
//colour = vec3(length(imageLoad(position_t, pix).rgb - imageLoad(geometry_prev_position, pix).rgb));
//const vec4 prev_colour = imageLoad(prev_dest, pix);
//colour = mix(colour, prev_colour.rgb, vec3(.9));
imageStore(out_dest, pix, vec4(colour, depth));
//imageStore(out_dest, pix, imageLoad(light_poly, pix));
imageStore(out_dest, pix, vec4(colour, 0./*unused*/));
}