floyd steinberg dithering

floyd_steinberg.zig

const std = @import("std");
const pnmdec = @import("pnmdec.zig");

pub fn quantFloydSteinberg(
    rows: usize,
    cols: usize,
    out: [*]u1,
    in: [*]u8,
) void {
    for (0..rows) |row_cursor| {
        const down_row = @min(rows - 1, row_cursor + 1);
        for (0..cols) |col_cursor| {
            const right_col = @min(cols - 1, col_cursor + 1);
            const left_col = col_cursor -| 1;

            const in_sample = in[row_cursor * cols + col_cursor];

            // compute positions
            const right = &in[row_cursor * cols + right_col];
            const down_right = &in[down_row * cols + right_col];
            const down_center = &in[down_row * cols + col_cursor];
            const down_left = &in[down_row * cols + left_col];

            // quantize
            const out_sample = @intFromBool(in_sample > 127);
            out[row_cursor * cols + col_cursor] = out_sample;

            // dequantize and compute error
            const dequant_sample = @as(i16, out_sample) * 255;
            const quant_error = in_sample - dequant_sample;

            // compute diffused error
            const right_error = @divTrunc(quant_error * 7, 16);
            const down_center_error = @divTrunc(quant_error * 5, 16);
            const down_right_error = @divTrunc(quant_error * 1, 16);
            const down_left_error = @divTrunc(quant_error * 3, 16);

            // apply error and clamp
            down_right.* = @intCast(@max(
                std.math.minInt(u8),
                @min(std.math.maxInt(u8), down_right.* + down_right_error),
            ));
            right.* = @intCast(@max(
                std.math.minInt(u8),
                @min(std.math.maxInt(u8), right.* + right_error),
            )); // down_right is == right in the last row
            down_center.* = @intCast(@max(
                std.math.minInt(u8),
                @min(std.math.maxInt(u8), down_center.* + down_center_error),
            ));
            down_left.* = @intCast(@max(
                std.math.minInt(u8),
                @min(std.math.maxInt(u8), down_left.* + down_left_error),
            ));
        }
    }
}

pub fn main() !void {
    const allocator = std.heap.smp_allocator;

    const rd = std.io.getStdIn().reader();

    const h = try pnmdec.parseHeader(rd);
    if (h.format != .pgm or h.maxval != 255) return error.Not8BitPgm;

    const image = try allocator.alloc(u8, h.height * h.width);
    defer allocator.free(image);

    try rd.readNoEof(image);

    const out = try allocator.alloc(u1, h.height * h.width);
    defer allocator.free(out);

    quantFloydSteinberg(h.height, h.width, out.ptr, image.ptr);

    try std.io.getStdOut().writer().print("P1\n{d} {d}\n", .{ h.width, h.height });

    for (0..h.height) |row_cursor| {
        for (0..h.width - 1) |col_cursor| {
            try std.io.getStdOut().writer().print("{d} ", .{1 ^ out[row_cursor * h.width + col_cursor]});
        }
        try std.io.getStdOut().writer().print("{d}\n", .{1 ^ out[row_cursor * h.width + (h.width - 1)]});
    }
}

pnmdec.zig

const std = @import("std");

pub const PnmFormat = enum(u8) {
    plain_pbm = '1',
    plain_pgm = '2',
    plain_ppm = '3',
    pbm = '4',
    pgm = '5',
    ppm = '6',

    fn fromVersion(byte: u8) ?PnmFormat {
        return std.meta.intToEnum(PnmFormat, byte) catch null;
    }

    fn isPlain(pf: PnmFormat) bool {
        return @intFromEnum(pf) <= '3';
    }

    fn hasMaxval(pf: PnmFormat) bool {
        return switch (pf) {
            .plain_pbm, .pbm => false,
            .plain_pgm, .plain_ppm, .pgm, .ppm => true,
        };
    }
};

pub const Header = struct {
    format: PnmFormat,
    width: usize,
    height: usize,
    maxval: u16,
};

pub fn parseHeader(in: anytype) !Header {
    const format: PnmFormat = get_format: {
        const magic = try in.readBytesNoEof(3);
        const m_format = PnmFormat.fromVersion(magic[1]);
        if (magic[0] != 'P' or magic[2] != '\n' or m_format == null) return error.ShitMagic;
        break :get_format m_format.?;
    };

    var tmp_buf: [512]u8 = undefined;

    // Parse the header.

    var width: usize = undefined;
    var height: usize = undefined;
    var maxval: u16 = 1;

    var state: enum {
        width,
        height,
        maxval,
    } = .width;

    outer: while (true) {
        const line_len = try in.readUntilDelimiter(&tmp_buf, '\n');
        const line = tmp_buf[0..line_len.len];

        if (line.len == 0) return error.EndOfStream;

        if (line[0] == '#') {
            continue;
        }

        var si = std.mem.tokenizeAny(u8, line, &std.ascii.whitespace);

        while (si.next()) |token| {
            switch (state) {
                .width => {
                    width = try std.fmt.parseUnsigned(usize, token, 10);
                    state = .height;
                },
                .height => {
                    height = try std.fmt.parseUnsigned(usize, token, 10);
                    state = if (format.hasMaxval()) .maxval else break :outer;
                },
                .maxval => {
                    maxval = try std.fmt.parseUnsigned(u16, token, 10);
                    break :outer;
                },
            }
        }
    }
    return .{
        .format = format,
        .width = width,
        .height = height,
        .maxval = maxval,
    };
}