Median and Rank Filters

const std = @import("std");
const canvas = @import("_canvas.zig");

const count = canvas.width * canvas.height;

var pixels: [count]u32 = undefined;
var scratch: [count]u32 = undefined;

fn sample(src: []const u32, x: i32, y: i32) u32 {
    const cx = std.math.clamp(x, 0, @as(i32, @intCast(canvas.width)) - 1);
    const cy = std.math.clamp(y, 0, @as(i32, @intCast(canvas.height)) - 1);
    return src[@as(usize, @intCast(cy)) * canvas.width + @as(usize, @intCast(cx))];
}

/// Salt and pepper: a fraction of pixels replaced outright by black or white.
/// This is the noise a bad sensor or a lossy channel produces, and it is the
/// case that separates the two filters here, because the corrupted values are
/// nothing like their neighbours.
///
/// A fixed multiplier and seed, so the picture is the same on every run and the
/// expected output can be checked in.
fn addSaltAndPepper(buf: []u32, per_mille: u32) void {
    var state: u64 = 0x2545F4914F6CDD1D;
    for (buf) |*p| {
        state = state *% 6364136223846793005 +% 1442695040888963407;
        const roll: u32 = @intCast((state >> 33) % 1000);
        if (roll < per_mille) {
            p.* = if (roll % 2 == 0) canvas.rgb(0, 0, 0) else canvas.rgb(255, 255, 255);
        }
    }
}

/// A 3x3 gaussian, for comparison. Every neighbour contributes something, which
/// is exactly the problem: a single white pixel is not removed, it is spread
/// over nine.
fn blur(dst: []u32, src: []const u32) void {
    const weights = [9]i32{ 1, 2, 1, 2, 4, 2, 1, 2, 1 };
    for (0..canvas.height) |y| {
        for (0..canvas.width) |x| {
            var r: i32 = 0;
            var g: i32 = 0;
            var b: i32 = 0;
            for (0..3) |ky| {
                for (0..3) |kx| {
                    const w = weights[ky * 3 + kx];
                    const c = canvas.channels(sample(
                        src,
                        @as(i32, @intCast(x + kx)) - 1,
                        @as(i32, @intCast(y + ky)) - 1,
                    ));
                    r += w * c.r;
                    g += w * c.g;
                    b += w * c.b;
                }
            }
            dst[y * canvas.width + x] = canvas.rgb(
                @intCast(@divTrunc(r, 16)),
                @intCast(@divTrunc(g, 16)),
                @intCast(@divTrunc(b, 16)),
            );
        }
    }
}

/// Sort the nine neighbours by luma and take the one at `rank`.
///
/// The whole family in one function. Rank 4 is the median, rank 0 the minimum
/// and rank 8 the maximum, and `rank` being `comptime` means each instantiation
/// is a separate specialised function rather than a branch in a hot loop.
///
/// The pixel is carried through whole rather than sorted per channel: taking
/// the median of each channel separately can produce a colour that appears
/// nowhere in the neighbourhood.
fn rankFilter(dst: []u32, src: []const u32, comptime rank: usize) void {
    comptime std.debug.assert(rank < 9);
    for (0..canvas.height) |y| {
        for (0..canvas.width) |x| {
            var window: [9]u32 = undefined;
            for (0..3) |ky| {
                for (0..3) |kx| {
                    window[ky * 3 + kx] = sample(
                        src,
                        @as(i32, @intCast(x + kx)) - 1,
                        @as(i32, @intCast(y + ky)) - 1,
                    );
                }
            }
            // Insertion sort. Nine elements is small enough that the O(n^2)
            // bound is beaten in practice by having no branches to mispredict
            // and no call overhead; production code uses a fixed sorting
            // network here, which is straight-line code with no loop at all.
            var i: usize = 1;
            while (i < window.len) : (i += 1) {
                const key = window[i];
                var j = i;
                while (j > 0 and canvas.luma(window[j - 1]) > canvas.luma(key)) : (j -= 1) {
                    window[j] = window[j - 1];
                }
                window[j] = key;
            }
            dst[y * canvas.width + x] = window[rank];
        }
    }
}

/// Count pixels that are pure black or pure white, which after the noise pass
/// is very nearly a count of the surviving noise.
fn countExtremes(buf: []const u32) struct { black: usize, white: usize } {
    var black: usize = 0;
    var white: usize = 0;
    for (buf) |p| {
        const l = canvas.luma(p);
        if (l == 0) black += 1;
        if (l >= 254) white += 1;
    }
    return .{ .black = black, .white = white };
}

/// Mean absolute luma difference between two pictures, in 1/100ths so the
/// integer print keeps two decimals.
fn meanError(a: []const u32, b: []const u32) u32 {
    var total: u64 = 0;
    for (a, b) |pa, pb| {
        const la: i32 = @intCast(canvas.luma(pa));
        const lb: i32 = @intCast(canvas.luma(pb));
        total += @abs(la - lb);
    }
    return @intCast((total * 100) / a.len);
}

fn printRow(out: *std.Io.Writer, label: []const u8, black: usize, white: usize, err: u32) !void {
    try out.print("{s:>10}: {d:>5} {d:>5} {d:>8}.{d:0>2}\n", .{ label, black, white, err / 100, err % 100 });
}

fn lumaRow(out: *std.Io.Writer, label: []const u8, buf: []const u32) !void {
    try out.print("{s:>9} |", .{label});
    for (5..13) |x| try out.print(" {d:>4}", .{canvas.luma(buf[10 * canvas.width + x])});
    try out.writeByte('\n');
}

pub fn main(init: std.process.Init) !void {
    var buf: [32768]u8 = undefined;
    var file_writer = std.Io.File.stdout().writerStreaming(init.io, &buf);
    const out = &file_writer.interface;

    canvas.scene(&pixels);
    addSaltAndPepper(&pixels, 60); // 6% of pixels
    try out.writeAll("6% salt and pepper:\n");
    try canvas.dump(out, &pixels);

    rankFilter(&scratch, &pixels, 4);
    try out.writeAll("\nmedian 3x3:\n");
    try canvas.dump(out, &scratch);

    const noisy = countExtremes(&pixels);
    const medianed = countExtremes(&scratch);
    blur(&scratch, &pixels);
    const blurred = countExtremes(&scratch);

    // Two measurements, because the first one lies. Counting pure black and
    // pure white says the gaussian removed every noisy pixel, and it did: it
    // averaged each one into a grey smudge nine pixels wide. Mean error against
    // the clean picture is the measurement that says whether the noise is gone
    // or merely spread out.
    var clean: [count]u32 = undefined;
    canvas.scene(&clean);

    rankFilter(&scratch, &pixels, 4);
    const median_error = meanError(&scratch, &clean);
    blur(&scratch, &pixels);
    const blur_error = meanError(&scratch, &clean);

    try out.print("\n            black white  mean luma error vs clean\n", .{});
    try printRow(out, "noisy", noisy.black, noisy.white, meanError(&pixels, &clean));
    try printRow(out, "gaussian", blurred.black, blurred.white, blur_error);
    try printRow(out, "median", medianed.black, medianed.white, median_error);

    // The edge, before and after each filter. The gaussian softens it because
    // it averages across it; the median does not, because on either side of the
    // edge the majority of the window is still on that side.
    canvas.scene(&pixels);
    try out.writeAll("\nluma along y=10, x=5..12 (the block's left edge)\n     x    |");
    for (5..13) |x| try out.print(" {d:>4}", .{x});
    try out.writeByte('\n');
    try lumaRow(out, "original", &pixels);
    blur(&scratch, &pixels);
    try lumaRow(out, "gaussian", &scratch);
    rankFilter(&scratch, &pixels, 4);
    try lumaRow(out, "median", &scratch);

    // Rank 0 and rank 8 are morphology. On a bright shape, the minimum filter
    // shrinks it by a pixel and the maximum grows it: erosion and dilation,
    // from the same function with a different index.
    canvas.scene(&pixels);
    var bright_before: usize = 0;
    for (pixels) |p| {
        if (canvas.luma(p) > 200) bright_before += 1;
    }

    rankFilter(&scratch, &pixels, 0);
    var eroded: usize = 0;
    for (scratch) |p| {
        if (canvas.luma(p) > 200) eroded += 1;
    }

    rankFilter(&scratch, &pixels, 8);
    var dilated: usize = 0;
    for (scratch) |p| {
        if (canvas.luma(p) > 200) dilated += 1;
    }

    try out.print("\npixels brighter than 200: {d} before, {d} after erosion, {d} after dilation\n", .{
        bright_before, eroded, dilated,
    });

    try out.flush();
}