Basic Monte Carlo noise remover

filter.cpp

#include <iostream>
#include <fstream>
#include <memory>
#include <string>
#include <algorithm>
#include <cmath>

#include <png.h>

struct Image {
    Image() : width(0), height(0) {}
    Image(size_t width, size_t height)
        : width(width), height(height), pixels(new float[4 * width * height])
    {}

    const float* row(int i) const { return pixels.get() + 4 * i * width; }
    const float* pixel(int i, int j) const { return row(j) + 4 * i; }
    float* row(int i) { return pixels.get() + 4 * i * width; }
    float* pixel(int i, int j) { return row(j) + 4 * i; }

    std::unique_ptr<float[]> pixels;
    size_t width;
    size_t height;
};

static float luminance(const float* pix) {
    return pix[0] * 0.2126f + pix[1] * 0.7152f + pix[2] * 0.0722f;
}

static float clamp(float a, float b, float c) {
    return std::max(b, std::min(c, a));
}

static void read_from_stream(png_structp png_ptr, png_bytep data, png_size_t length) {
    png_voidp a = png_get_io_ptr(png_ptr);
    ((std::istream*)a)->read((char*)data, length);
}

static void png_write_to_stream(png_structp png_ptr, png_bytep data, png_size_t length) {
    png_voidp a = png_get_io_ptr(png_ptr);
    ((std::ostream*)a)->write((const char*)data, length);
}

static void png_flush_stream(png_structp) {
    // Nothing to do
}

bool load_png(const std::string& path, Image& img) {
    std::ifstream file(path, std::ifstream::binary);
    if (!file)
        return false;

    // Read signature
    char sig[8];
    file.read(sig, 8);
    if (!png_check_sig((unsigned char*)sig, 8))
        return false;

    png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (!png_ptr)
        return false;

    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        png_destroy_read_struct(&png_ptr, nullptr, nullptr);
        return false;
    }

    if (setjmp(png_jmpbuf(png_ptr))) {
        png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
        return false;
    }

    png_set_sig_bytes(png_ptr, 8);
    png_set_read_fn(png_ptr, (png_voidp)&file, read_from_stream);
    png_read_info(png_ptr, info_ptr);

    img.width    = png_get_image_width(png_ptr, info_ptr);
    img.height   = png_get_image_height(png_ptr, info_ptr);

    png_uint_32 color_type = png_get_color_type(png_ptr, info_ptr);
    png_uint_32 bit_depth  = png_get_bit_depth(png_ptr, info_ptr);

    // Expand paletted and grayscale images to RGB
    if (color_type == PNG_COLOR_TYPE_PALETTE) {
        png_set_palette_to_rgb(png_ptr);
    } else if (color_type == PNG_COLOR_TYPE_GRAY ||
               color_type == PNG_COLOR_TYPE_GRAY_ALPHA) {
        png_set_gray_to_rgb(png_ptr);
    }

    // Transform to 8 bit per channel
    if (bit_depth == 16)
        png_set_strip_16(png_ptr);

    // Get alpha channel when there is one
    if (png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS))
        png_set_tRNS_to_alpha(png_ptr);

    // Otherwise add an opaque alpha channel
    else
        png_set_filler(png_ptr, 0xFF, PNG_FILLER_AFTER);

    img.pixels.reset(new float[4 * img.width * img.height]);
    std::unique_ptr<png_byte[]> row_bytes(new png_byte[img.width * 4]);
    for (size_t y = 0; y < img.height; y++) {
        png_read_row(png_ptr, row_bytes.get(), nullptr);
        float* img_row = img.row(y);
        const float inv = 1.0f / 255.0f;
        for (size_t x = 0; x < img.width; x++) {
            img_row[x * 4 + 0] = row_bytes[x * 4 + 0] * inv;
            img_row[x * 4 + 1] = row_bytes[x * 4 + 1] * inv;
            img_row[x * 4 + 2] = row_bytes[x * 4 + 2] * inv;
            img_row[x * 4 + 3] = row_bytes[x * 4 + 3] * inv;
        }
    }

    png_destroy_read_struct(&png_ptr, &info_ptr, nullptr);
    return true;
}

bool save_png(const std::string& path, const Image& img) {
    std::ofstream file(path, std::ofstream::binary);
    if (!file)
        return false;

    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr, nullptr);
    if (!png_ptr)
        return false;

    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (!info_ptr) {
        png_destroy_read_struct(&png_ptr, nullptr, nullptr);
        return false;
    }

    std::unique_ptr<uint8_t[]> row(new uint8_t[img.width * 4]);
    if (setjmp(png_jmpbuf(png_ptr))) {
        png_destroy_write_struct(&png_ptr, &info_ptr);
        return false;
    }

    png_set_write_fn(png_ptr, &file, png_write_to_stream, png_flush_stream);

    png_set_IHDR(png_ptr, info_ptr, img.width, img.height,
                 8, PNG_COLOR_TYPE_RGBA, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    png_write_info(png_ptr, info_ptr);

    for (size_t y = 0; y < img.height; y++) {
        const float* input = img.row(y);
        for (size_t x = 0; x < img.width; x++) {
            row[x * 4 + 0] = clamp(input[4 * x + 0], 0.0f, 1.0f) * 255.0f;
            row[x * 4 + 1] = clamp(input[4 * x + 1], 0.0f, 1.0f) * 255.0f;
            row[x * 4 + 2] = clamp(input[4 * x + 2], 0.0f, 1.0f) * 255.0f;
            row[x * 4 + 3] = clamp(input[4 * x + 3], 0.0f, 1.0f) * 255.0f;
        }
        png_write_row(png_ptr, row.get());
    }

    png_write_end(png_ptr, info_ptr);
    png_destroy_write_struct(&png_ptr, &info_ptr);

    return true;
}

int main(int argc, char** argv) {
    if (argc < 3 || argc > 5) {
        std::cout << "usage: filter input.png output.png [threshold] [size]" << std::endl;
        return 1;
    }

    Image img;
    if (!load_png(argv[1], img)) {
        std::cerr << "Cannot load image" << std::endl;
        return 1;
    }

    int size = argc >= 5 ? strtol(argv[4], nullptr, 10) : 1;
    float threshold = argc >= 4 ? strtof(argv[3], nullptr) : 0.1f;
    int d = 2 * size + 1;
    int w = img.width, h = img.height;

    Image new_img(img.width, img.height);
    #pragma omp parallel
    {
        std::unique_ptr<int[]> ids(new int[d * d]);
        std::unique_ptr<float[]> values(new float[d * d]);

        #pragma omp for collapse(2)
        for (int y = 0; y < h; ++y) {
            for (int x = 0; x < w; ++x) {
                for (int i = -size, p = 0; i <= size; ++i) {
                    for (int j = -size; j <= size; ++j, p++) {
                        int xi = x + i < 0 || x + i >= w ? 0 : x + i;
                        int yj = y + j < 0 || y + j >= h ? 0 : y + j;
                        const float* pix = img.pixel(xi, yj);
                        values[p] = luminance(pix);
                        ids[p] = p;
                    }
                }

                std::sort(ids.get(), ids.get() + d * d, [&] (int a, int b) {
                    return values[a] < values[b];
                });

                int c = d * d / 2;
                int m = ids[c];
                auto out = new_img.pixel(x, y);
                if (std::fabs(values[m] - values[c]) > threshold * (values[ids[d * d - 1]] - values[ids[0]])) {
                    int xi = x + (m % d - size);
                    int yj = y + (m / d - size);
                    if (xi < 0 || xi >= w) xi = x;
                    if (yj < 0 || yj >= h) yj = y;
                    auto in = img.pixel(xi, yj);
                    std::copy(in, in + 4, out);
                } else {
                    auto in = img.pixel(x, y);
                    std::copy(in, in + 4, out);
                }
            }
        }
    }

    save_png(argv[2], new_img);
    return 0;
}