ray-tracing/denoiser.cpp at main · GlitchedCode922/ray-tracing · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
#include "denoiser.h"
#include "vector.h"
#include <algorithm>
#include <vector>

void Denoiser::medianFilter(std::vector<Vector3>& image, int width, int height) {
    std::vector<Vector3> temp_image = image;
    int kernel_size = 3;
    int k = kernel_size / 2;
    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            std::vector<double> r_values;
            std::vector<double> g_values;
            std::vector<double> b_values;
            for (int ky = -k; ky <= k; ky++) {
                for (int kx = -k; kx <= k; kx++) {
                    int nx = std::min(std::max(x + kx, 0), width - 1);
                    int ny = std::min(std::max(y + ky, 0), height - 1);
                    Vector3 pixel = temp_image[ny * width + nx];
                    r_values.push_back(pixel.x);
                    g_values.push_back(pixel.y);
                    b_values.push_back(pixel.z);
                }
            }
            std::sort(r_values.begin(), r_values.end());
            std::sort(g_values.begin(), g_values.end());
            std::sort(b_values.begin(), b_values.end());
            int median_index = r_values.size() / 2;
            image[y * width + x] = Vector3(
                r_values[median_index],
                g_values[median_index],
                b_values[median_index]
            );
        }
    }
}

void Denoiser::bilateralFilter(std::vector<Vector3> &image, int width, int height, double sigma_spatial, double sigma_color) {
    std::vector<Vector3> temp_image = image;
    int kernel_size = 5;
    int k = kernel_size / 2;
    double two_sigma_spatial_sq = 2 * sigma_spatial * sigma_spatial;
    double two_sigma_color_sq = 2 * sigma_color * sigma_color;

    for (int y = 0; y < height; y++) {
        for (int x = 0; x < width; x++) {
            Vector3 sum(0, 0, 0);
            double weight_sum = 0.0;
            Vector3 center_pixel = temp_image[y * width + x];

            for (int ky = -k; ky <= k; ky++) {
                for (int kx = -k; kx <= k; kx++) {
                    int nx = std::min(std::max(x + kx, 0), width - 1);
                    int ny = std::min(std::max(y + ky, 0), height - 1);
                    Vector3 neighbor_pixel = temp_image[ny * width + nx];

                    double spatial_dist_sq = kx * kx + ky * ky;
                    double color_dist_sq = (center_pixel - neighbor_pixel).magnitude();
                    color_dist_sq *= color_dist_sq;

                    double spatial_weight = exp(-spatial_dist_sq / two_sigma_spatial_sq);
                    double color_weight = exp(-color_dist_sq / two_sigma_color_sq);
                    double weight = spatial_weight * color_weight;

                    sum += neighbor_pixel * weight;
                    weight_sum += weight;
                }
            }
            image[y * width + x] = sum / weight_sum;
        }
    }
}

void Denoiser::atrousWavelet(std::vector<Vector3>& image, int width, int height, int iterations, double sigma_color) {
    if (iterations <= 0) return;
    std::vector<Vector3> src = image;
    std::vector<Vector3> tmp(image.size());

    // 3-tap binomial kernel weights
    const double kernel[3] = {1.0/4.0, 2.0/4.0, 1.0/4.0};
    const double two_sigma_color_sq = 2.0 * sigma_color * sigma_color;

    for (int iter = 0; iter < iterations; ++iter) {
        int step = 1 << iter; // dilation: 1,2,4,...

        // convolve separably: horizontal then vertical (with dilation)
        // Horizontal pass
        for (int y = 0; y < height; ++y) {
            for (int x = 0; x < width; ++x) {
                Vector3 center = src[y * width + x];
                Vector3 sum(0,0,0);
                double wsum = 0.0;

                for (int k = -1; k <= 1; ++k) {
                    int sx = x + k * step;
                    if (sx < 0) sx = 0;
                    if (sx >= width) sx = width - 1;
                    Vector3 sample = src[y * width + sx];

                    // color distance in RGB space
                    Vector3 diff = center - sample;
                    double dist2 = diff.x*diff.x + diff.y*diff.y + diff.z*diff.z;
                    double color_w = exp(-dist2 / two_sigma_color_sq);
                    double g = kernel[k + 1] * color_w;

                    sum += sample * g;
                    wsum += g;
                }
                tmp[y * width + x] = sum / (wsum > 0.0 ? wsum : 1.0);
            }
        }

        // Vertical pass (write back into src for next iteration)
        for (int y = 0; y < height; ++y) {
            for (int x = 0; x < width; ++x) {
                Vector3 center = tmp[y * width + x];
                Vector3 sum(0,0,0);
                double wsum = 0.0;

                for (int k = -1; k <= 1; ++k) {
                    int sy = y + k * step;
                    if (sy < 0) sy = 0;
                    if (sy >= height) sy = height - 1;
                    Vector3 sample = tmp[sy * width + x];

                    Vector3 diff = center - sample;
                    double dist2 = diff.x*diff.x + diff.y*diff.y + diff.z*diff.z;
                    double color_w = exp(-dist2 / two_sigma_color_sq);
                    double g = kernel[k + 1] * color_w;

                    sum += sample * g;
                    wsum += g;
                }
                src[y * width + x] = sum / (wsum > 0.0 ? wsum : 1.0);
            }
        }
    }

    // write back result
    image = src;
}