Lens filter using three separable complex gaussian kernels

lens.c

#include <complex.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/fcntl.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <string.h>
#include <math.h>

/*
  Le fichier d'entrée doit être un buffer brut RGB
 */
float *read_img(const char *n) {
    struct stat fs;
    int fd = open(n, O_RDONLY);
    fstat(fd, &fs);

    unsigned char *filecontent = mmap(0, fs.st_size, PROT_READ, MAP_PRIVATE, fd, 0);
    float *ret = malloc(sizeof(float) * fs.st_size);
    for (unsigned i = 0; i < fs.st_size; ++i) {
        ret[i] = (float)filecontent[i] / 255.0f;
    }
    munmap(filecontent, fs.st_size);
    close(fd);
    return ret;
}

void write_img(float *data, unsigned n, const char *fn) {
    int fd = open(fn, O_RDWR | O_CREAT | O_TRUNC, 0644);
    unsigned char *buf = malloc(n * 3);
    for (unsigned i = 0; i < n * 3; ++i) {
        buf[i] = data[i] * 255;
    }
    write(fd, buf, n * 3);
    close(fd);
    free(buf);
}

void build_kernel(complex float *k, int r, float scale, float a, float b) {
    unsigned w = r * 2 + 1;
    complex float t[w];
    for (int i = -r; i < r + 1; ++i) {
        t[i + r] = (i * scale) / r;
    }
    for (int i = 0; i < w; ++i) {
        // Gaussien reel/imag
        k[i] =
            exp(-a * pow(t[i], 2)) * cos(b * pow(t[i], 2)) +
            exp(-a * pow(t[i], 2)) * sin(b * pow(t[i], 2)) * I;
    }
}

void apply_gamma(float *data, unsigned n, float g) {
    n *= 3;
    for (unsigned i = 0; i < n; ++i)
        data[i] = pow(data[i], g);
}

complex float *make_complex(float *data, unsigned n) {
    complex float *ret = malloc(3 * sizeof(complex float) * n);
    for (unsigned i = 0; i < 3 * n; ++i)
        ret[i] = data[i] + 0I;
    return ret;
}

void clamp_image(float *data, unsigned n, float min, float max) {
    n *= 3;
    for (unsigned i = 0; i < n; ++i) {
        data[i] = data[i] < min ? min : (data[i] > max ? max : data[i]);
    }
}

void make_real(float *out, complex float *in, int w, int h, float a, float b) {
    unsigned n = w * h * 3;
    for (unsigned i = 0; i < n; ++i) {
        out[i] = a * creal(in[i]) + b * cimag(in[i]);
    }   
}

void image_acc(float *out, float *in, unsigned n) {
    n *= 3;
    for (unsigned i = 0; i < n; ++i) {
        out[i] += in[i];
    }
}

#define PIXEL_AT(w, h, x, y) (3 * (((y) * (w)) + (x)))

void hpass(complex float *m, int w, int h,
           complex float *kernel, int kl,
           complex float *out) {
    const int offset = kl / 2; 
    for (int y = 0; y < h; ++y)
        for (int x = offset; x < w - offset; ++x) {
            // pour chaque rgb
            /* TODO: Inverser cette boucle */
            for (int c = 0; c < 3; ++c) {
                complex float acc = 0;
                for (int i = 0; i < kl; ++i) {
                    unsigned o = PIXEL_AT(w, h, x - offset + i, y) + c;
                    acc += kernel[i] * m[o];
                }
                unsigned o = PIXEL_AT(w - (kl - 1), h, x - offset, y) + c;
                out[o] = acc;
            }
        }
}

// Convolution par la transposée
void vpass(complex float *m, int w, int h,
           complex float *kernel, int kl,
           complex float *out) {
    const int offset = kl / 2; 
    for (int y = offset; y < h - offset; ++y)
        for (int x = 0; x < w; ++x) {
            // pour chaque rgb
            /* TODO: Inverser cette boucle */
            for (int c = 0; c < 3; ++c) {
                complex float acc = 0;
                for (int i = 0; i < kl; ++i) {
                    unsigned o = PIXEL_AT(w, h, x, y - offset + i) + c;
                    acc += kernel[i] * m[o];
                }
                unsigned o = PIXEL_AT(w, h - (kl - 1), x, y - offset) + c;
                out[o] = acc;
            }
        }
}

void norm_kernels(complex float *kernels, int n, int r, float (*coef_table)[4]) {
    float acc = 0;

    for (int k = 0; k < n; ++k) {
        complex float *kp = kernels + k * (r * 2 + 1);
        float *p = &coef_table[k][2];
        for (int i = 0; i < r * 2 + 1; ++i)
            for (int j = 0; j < r * 2 + 1; ++j)
                // Produit intérieur
                acc +=
                    p[0] * (creal(kp[i]) * creal(kp[j]) - cimag(kp[i]) * cimag(kp[j])) +
                    p[1] * (creal(kp[i]) * cimag(kp[j]) + cimag(kp[i]) * creal(kp[j]));
    }
    float inv = 1.0f / sqrt(acc);
    int t = (r * 2 + 1) * n;
    for (int k = 0; k < t; ++k) {
        kernels[k] *= inv;
    }
}

int main(int argc, char *argv[argc]) {
    int w = 128; // Dimensions de l'image
    int h = 128;
    int r = 20; // Rayon du noyau

    int ow = w - r * 2;
    int oh = h - r * 2;

    float *img = read_img(argv[1]);
    float *output = calloc(sizeof(float) * ow, oh * 3);

    float three_coefs[3][4] = {
        /* Amplitude R et I | Poids R et I */
        { 2.17649,  5.043495,  1.621035, -2.105439 },
        { 1.019306, 9.027613, -0.28086,  -0.162882 },
        { 2.81511,  1.597273, -0.366471, 10.300301 }
    };
    float s = 1.2;

    complex float *kernels = malloc(sizeof(complex float) * 3 * (r * 2 + 1));
    for (int i = 0; i < 3; ++i) {
        build_kernel(kernels + i * ((r * 2 + 1)), r, s, three_coefs[i][0], three_coefs[i][1]);
    }

    norm_kernels(kernels, 3, r, three_coefs);

    apply_gamma(img, w * h, 3.0);
    complex float *cimg = make_complex(img, w * h);

    /* 
       Chaque convolution de chaque composant 
       produit une image complexe.
       Les composant reel et imaginaires sont remélangé selon les 
       poids en une valeure reelle.
       Je considere que la convolution cause un troncage
     */
    float *buff = malloc(3 * sizeof(float) * ow * oh);
    // passe horizontale, détruit r colonnes
    complex float *scratch1 = malloc(3 * sizeof(complex float) * ow * h);
    // passe verticale, détruit r lignes
    complex float *scratch2 = malloc(3 * sizeof(complex float) * ow * oh);
    for (int i = 0; i < 3; ++i) {
        float *params = &three_coefs[i][2];
        complex float *kernel = kernels + i * (r * 2 + 1);
        hpass(cimg, w, h, kernel, r * 2 + 1, scratch1);
        vpass(scratch1, ow, h, kernel, r * 2 + 1, scratch2);
        make_real(buff, scratch2, ow, oh, params[0], params[1]);
        image_acc(output, buff, ow * oh);
    }
    
    apply_gamma(output, ow * oh, 1.0 / 3.0);
    clamp_image(output, ow * oh, 0, 1);

    write_img(output, ow * oh, argv[2]);
    free(img);
    free(cimg);
    free(scratch1);
    free(scratch2);
    free(buff);
    free(output);
    free(kernels);
    return 0;
}