Skip to main content

KERNEL_HISTOGRAM_EQUALIZE

Constant KERNEL_HISTOGRAM_EQUALIZE 

Source
pub const KERNEL_HISTOGRAM_EQUALIZE: &str = r#"
// Histogram computation and equalization.
// Pass 1: Compute histogram (atomically increment bins).
// Pass 2: Use CDF to remap values.
// Selected by PASS_MODE define: 0 = histogram, 1 = CDF prefix sum, 2 = equalize.

layout(local_size_x = 256) in;

#ifndef PASS_MODE
#define PASS_MODE 0
#endif

#ifndef BIN_COUNT
#define BIN_COUNT 256
#endif

layout(std430, binding = 0) buffer InputData {
    float input_data[];
};

layout(std430, binding = 1) buffer Histogram {
    uint histogram[];
};

layout(std430, binding = 2) buffer CDF {
    float cdf[];
};

layout(std430, binding = 3) buffer OutputData {
    float output_data[];
};

uniform uint u_element_count;
uniform float u_min_value;
uniform float u_max_value;

// Shared memory for local histogram accumulation
shared uint local_hist[BIN_COUNT];

void main() {
    uint idx = gl_GlobalInvocationID.x;
    uint lid = gl_LocalInvocationID.x;

#if PASS_MODE == 0
    // Pass 0: Build histogram
    // Initialize shared histogram
    if (lid < uint(BIN_COUNT)) {
        local_hist[lid] = 0u;
    }
    barrier();

    if (idx < u_element_count) {
        float val = input_data[idx];
        float norm = clamp((val - u_min_value) / (u_max_value - u_min_value), 0.0, 1.0);
        uint bin = min(uint(norm * float(BIN_COUNT - 1)), uint(BIN_COUNT - 1));
        atomicAdd(local_hist[bin], 1u);
    }
    barrier();

    // Merge local histogram into global
    if (lid < uint(BIN_COUNT)) {
        atomicAdd(histogram[lid], local_hist[lid]);
    }

#elif PASS_MODE == 1
    // Pass 1: Build CDF from histogram (single workgroup, sequential for simplicity)
    if (idx == 0u) {
        uint running = 0u;
        for (uint i = 0u; i < uint(BIN_COUNT); i++) {
            running += histogram[i];
            cdf[i] = float(running) / float(u_element_count);
        }
    }

#elif PASS_MODE == 2
    // Pass 2: Apply equalization using CDF
    if (idx < u_element_count) {
        float val = input_data[idx];
        float norm = clamp((val - u_min_value) / (u_max_value - u_min_value), 0.0, 1.0);
        uint bin = min(uint(norm * float(BIN_COUNT - 1)), uint(BIN_COUNT - 1));
        float equalized = cdf[bin];
        output_data[idx] = equalized * (u_max_value - u_min_value) + u_min_value;
    }

#endif
}
"#;
Expand description

Histogram equalization kernel (two passes: histogram + equalize).