@woosh/meep-engine/src/engine/graphics/shaders/DenoiseShader.js

Pure JavaScript game engine. Fully featured and production ready.

import { Vector2 } from "three";
import { pdf_normal } from "../../../core/math/physics/pdf/pdf_normal.js";

const KERNEL_SIZE = 15;

const fragment = `
// bilateral filter, based on https://www.shadertoy.com/view/4dfGDH# and
// http://people.csail.mit.edu/sparis/bf_course/course_notes.pdf
// A bilateral filter is a non-linear, edge-preserving, and noise-reducing smoothing filter for images.
// It replaces the intensity of each pixel with a weighted average of intensity values from nearby pixels.
// This weight can be based on a Gaussian distribution. Crucially, the weights depend not only on
// Euclidean distance of pixels, but also on the radiometric differences (e.g., range differences, such
// as color intensity, depth distance, etc.). This preserves sharp edges.

float normpdf3(in vec3 v, in float sigma) {
    return 0.39894 * exp(-0.5 * dot(v, v) / (sigma * sigma)) / sigma;
}

// filter size
#define MSIZE ${KERNEL_SIZE}
varying vec2 vUv;
uniform sampler2D source;
uniform vec2 pixelOffset;
uniform vec2 sigmas;
uniform float bZnorm;
uniform float kernel[MSIZE];

void main(void) {

    // range sigma - controls blurriness based on a pixel distance
    float sigma = sigmas.x;
    
    // domain sigma - controls blurriness based on a pixel similarity (to preserve edges)
    float bSigma = sigmas.y;
    vec4 pixel = texture2D(source, vUv);
    vec4 accumulatedColor = vec4(0.0);
    float accumulatedFactor = 0.0;
    
    // read out the texels
    const int kSize = (MSIZE-1)/2;
    
    for (int i = -kSize; i <= kSize; ++i) {
        for (int j = -kSize; j <= kSize; ++j) {

            // sample the pixel with offset
            vec2 coord = vUv + vec2(float(i), float(j)) * pixelOffset;
            vec4 rgba = texture2D(source, coord);
       
            // bilateral factors
            float factor = kernel[kSize + j] * kernel[kSize + i];
            factor *= normpdf3(rgba.rgb - pixel.rgb, bSigma) * bZnorm;
            
            // accumulate
            accumulatedColor += factor * rgba;
            accumulatedFactor += factor;
        }
    }
    
    gl_FragColor = accumulatedColor / accumulatedFactor;
}
`;

/**
 *
 * @param {number} sigma
 * @param {number} filterSmoothness
 * @constructor
 */
export function DenoiseShader(sigma = 10, filterSmoothness = 0.2) {
    // compute zNorm
    const zNorm = 1 / pdf_normal(0, filterSmoothness);

    // build kernel
    const kernel = new Float32Array(KERNEL_SIZE);
    const kSize = Math.floor((KERNEL_SIZE - 1) / 2);
    for (let i = 0; i <= kSize; i++) {
        const v = pdf_normal(i, sigma);

        kernel[kSize + i] = v;
        kernel[kSize - i] = v;
    }

    return {
        uniforms: {
            source: {
                type: 't',
                value: null
            },
            pixelOffset: {
                type: 'v2',
                /**
                 * Holds single pixel size in UV space ( 1 / texture_size )
                 */
                value: new Vector2(1, 1)
            },
            sigmas: {
                type: "v2", value: new Vector2(sigma, sigma)
            },
            bZnorm: {
                type: 'f',
                value: zNorm
            },
            kernel: {
                type: 'f',
                value: kernel
            }
        },
        fragmentShader: fragment
    };
}