Revision control
++ // Copyright (C) 2025 The Qt Company Ltd.
++ // SPDX-License-Identifier: LicenseRef-Qt-Commercial OR GPL-3.0-only
++
++ // NL-Means filter (NLM)
++ //
++ // Based on:
++ // Antoni Buades, Bartomeu Coll, and Jean-Michel Morel,
++ // Non-Local Means Denoising, Image Processing On Line, 1 (2011), pp. 208–212.
++ //
++ // The notable difference in this implementation is the addition of a mask that
++ // makes sure that only pixels in the same uv chart are used.
++ //
++ // The inputData[] is a f32 buffer and outputData[] is a u32 buffer to be
++ // able to use atomicAdd for accumulation, which is then divded by
++ // pixelCounts[]. To handle u32 all values are scaled by 1000.f
++
++ #version 430 core
++ layout (local_size_x = 8, local_size_y = 8, local_size_z = 1) in;
++
++ layout(std140, binding = 0) uniform UniformBuffer
++ {
++ float sigma;
++ float width; // int
++ float height; // int
++ } ubuf;
++
++ layout(std430, binding = 1) buffer InputBuffer {
++ float inputData[]; // Packed as: [y][x][rgb]
++ };
++
++ layout(rgba8, binding = 2) uniform image2D imgMask;
++
++ layout(std430, binding = 3) buffer OutputBuffer {
++ uint outputData[]; // Using u32 for atomic add
++ };
++
++ layout(std430, binding = 4) buffer CountBuffer {
++ uint pixelCounts[];
++ };
++
++ bool isInsideImage(ivec2 coord) {
++ return coord.x >= 0 && coord.x < int(ubuf.width) &&
++ coord.y >= 0 && coord.y < int(ubuf.height);
++ }
++
++ bool isSameChart(ivec2 coord, vec4 referenceColor) {
++ return imageLoad(imgMask, coord) == referenceColor;
++ }
++
++ vec3 getPixel(ivec2 coord) {
++ int width = int(ubuf.width);
++ int pixelIndex = 3 * (coord.y * width + coord.x);
++ return vec3(inputData[pixelIndex], inputData[pixelIndex + 1], inputData[pixelIndex + 2]);
++ }
++
++ const float scale = 1000.0;
++
++ void main() {
++ const int x = int(gl_GlobalInvocationID.x);
++ const int y = int(gl_GlobalInvocationID.y);
++ const vec4 chartColor = imageLoad(imgMask, ivec2(x, y));
++
++ // If this is not a lightmapped pixel, just return
++ if (chartColor == vec4(0, 0, 0, 0)) {
++ return;
++ }
++
++ float sigma = clamp(ubuf.sigma, 0.0f, 100.0f);
++ int searchWindowRadius;
++ int searchBlockRadius;
++ float filterMultiplier;
++
++ const int imgWidth = int(ubuf.width);
++ const int imgHeight = int(ubuf.height);
++
++ // Set window size and block radius based on sigma
++ if (sigma <= 25.0f) {
++ searchWindowRadius = 1;
++ searchBlockRadius = 10;
++ filterMultiplier = 0.55f;
++ } else if (sigma <= 55.0f) {
++ searchWindowRadius = 2;
++ searchBlockRadius = 17;
++ filterMultiplier = 0.4f;
++ } else {
++ searchWindowRadius = 3;
++ searchBlockRadius = 17;
++ filterMultiplier = 0.35f;
++ }
++
++ float sigmaSquared = sigma * sigma;
++ float h = filterMultiplier * sigma;
++ float hSquared = h * h;
++
++ int patchRadius = min(searchWindowRadius, min(min(x, y), min(imgWidth - 1 - x, imgHeight - 1 - y)));
++
++ int searchXStart = max(x - searchBlockRadius, patchRadius);
++ int searchXEnd = min(x + searchBlockRadius, imgWidth - 1 - patchRadius);
++ int searchYStart = max(y - searchBlockRadius, patchRadius);
++ int searchYEnd = min(y + searchBlockRadius, imgHeight - 1 - patchRadius);
++
++ int patchSize = 2 * patchRadius + 1;
++ int patchArea = patchSize * patchSize;
++
++ vec3 patchAccum[121]; // Max patch size 11x11 = 121
++ for (int i = 0; i < 121; i++) {
++ patchAccum[i] = vec3(0.0);
++ }
++
++ float maxWeight = 0.0f;
++ float totalWeight = 0.0f;
++
++ for (int sy = searchYStart; sy <= searchYEnd; sy++) {
++ for (int sx = searchXStart; sx <= searchXEnd; sx++) {
++ if (sx == x && sy == y) continue;
++ if (!isSameChart(ivec2(sx, sy), chartColor)) continue;
++
++ float distance = 0.0f;
++
++ for (int dy = -patchRadius; dy <= patchRadius; dy++) {
++ for (int dx = -patchRadius; dx <= patchRadius; dx++) {
++ ivec2 refCoord = ivec2(x + dx, y + dy);
++ ivec2 sampleCoord = ivec2(sx + dx, sy + dy);
++
++ if (!isSameChart(refCoord, chartColor) ||
++ !isSameChart(sampleCoord, chartColor)) {
++ continue;
++ }
++
++ vec3 refPixel = getPixel(refCoord);
++ vec3 samplePixel = getPixel(sampleCoord);
++ distance += dot(refPixel - samplePixel, refPixel - samplePixel);
++ }
++ }
++
++ distance = max(distance - 2.0f * 3 * patchArea * sigmaSquared, 0.0f);
++ float weight = exp(-distance / hSquared);
++
++ maxWeight = max(maxWeight, weight);
++ totalWeight += weight;
++
++ for (int dy = -patchRadius; dy <= patchRadius; dy++) {
++ for (int dx = -patchRadius; dx <= patchRadius; dx++) {
++ ivec2 coord = ivec2(sx + dx, sy + dy);
++ if (!isSameChart(coord, chartColor)) continue;
++
++ vec3 color = getPixel(coord);
++ int px = dx + patchRadius;
++ int py = dy + patchRadius;
++ int patchIndex = py * patchSize + px;
++
++ patchAccum[patchIndex] += weight * color;
++ }
++ }
++ }
++ }
++
++ // Add reference patch itself with maxWeight
++ for (int dy = -patchRadius; dy <= patchRadius; dy++) {
++ for (int dx = -patchRadius; dx <= patchRadius; dx++) {
++ vec3 refColor = getPixel(ivec2(x + dx, y + dy));
++ int px = dx + patchRadius;
++ int py = dy + patchRadius;
++ int patchIndex = py * patchSize + px;
++ patchAccum[patchIndex] += maxWeight * refColor;
++ }
++ }
++
++ totalWeight += maxWeight;
++ const float epsilon = 1e-8f;
++
++ if (totalWeight <= epsilon)
++ return;
++
++ for (int dy = -patchRadius; dy <= patchRadius; dy++) {
++ for (int dx = -patchRadius; dx <= patchRadius; dx++) {
++ int outX = x + dx;
++ int outY = y + dy;
++ int px = dx + patchRadius;
++ int py = dy + patchRadius;
++ int patchIndex = py * patchSize + px;
++ vec3 finalColor = patchAccum[patchIndex] / totalWeight;
++ int pixelIndex = outY * imgWidth + outX;
++ int outputIndex = pixelIndex * 3;
++ atomicAdd(pixelCounts[pixelIndex], 1);
++ atomicAdd(outputData[outputIndex], uint(finalColor.r * scale));
++ atomicAdd(outputData[outputIndex + 1], uint(finalColor.g * scale));
++ atomicAdd(outputData[outputIndex + 2], uint(finalColor.b * scale));
++ }
++ }
++ }