Add an implementation of sharpening by FIR Wiener filters.
[movit] / deconvolution_sharpen_effect.cpp
1 // NOTE: Throughout, we use the symbol ⊙ for convolution.
2 // Since all of our signals are symmetrical, discrete correlation and convolution
3 // is the same operation, and so we won't make a difference in notation.
4
5
6 #include <math.h>
7 #include <assert.h>
8 #include <Eigen/Dense>
9 #include <Eigen/Cholesky>
10
11 #include "deconvolution_sharpen_effect.h"
12 #include "util.h"
13 #include "opengl.h"
14
15 using namespace Eigen;
16
17 DeconvolutionSharpenEffect::DeconvolutionSharpenEffect()
18         : R(5),
19           circle_radius(2.0f),
20           gaussian_radius(0.0f),
21           correlation(0.95f),
22           noise(0.01f)
23 {
24         register_int("matrix_size", &R);
25         register_float("circle_radius", &circle_radius);
26         register_float("gaussian_radius", &gaussian_radius);
27         register_float("correlation", &correlation);
28         register_float("noise", &noise);
29 }
30
31 std::string DeconvolutionSharpenEffect::output_fragment_shader()
32 {
33         char buf[256];
34         sprintf(buf, "#define R %u\n", R);
35         return buf + read_file("deconvolution_sharpen_effect.frag");
36 }
37
38 namespace {
39
40 // Integral of sqrt(r² - x²) dx over x=0..a.
41 float circle_integral(float a, float r)
42 {
43         assert(a >= 0.0f);
44         if (a <= 0.0f) {
45                 return 0.0f;
46         }
47         if (a >= r) {
48                 return 0.25f * M_PI * r * r;
49         }
50         return 0.5f * (a * sqrt(r*r - a*a) + r*r * asin(a / r));
51 }
52
53 // Yields the impulse response of a circular blur with radius r.
54 // We basically look at each element as a square centered around (x,y),
55 // and figure out how much of its area is covered by the circle.
56 float circle_impulse_response(int x, int y, float r)
57 {
58         if (r < 1e-3) {
59                 // Degenerate case: radius = 0 yields the impulse response.
60                 return (x == 0 && y == 0) ? 1.0f : 0.0f;
61         }
62
63         // Find the extents of this cell. Due to symmetry, we can cheat a bit
64         // and pretend we're always in the upper-right quadrant, except when
65         // we're right at an axis crossing (x = 0 or y = 0), in which case we
66         // simply use the evenness of the function; shrink the cell, make
67         // the calculation, and down below we'll normalize by the cell's area.
68         float min_x, max_x, min_y, max_y;
69         if (x == 0) {
70                 min_x = 0.0f;
71                 max_x = 0.5f;
72         } else {
73                 min_x = abs(x) - 0.5f;
74                 max_x = abs(x) + 0.5f;
75         }
76         if (y == 0) {
77                 min_y = 0.0f;
78                 max_y = 0.5f;
79         } else {
80                 min_y = abs(y) - 0.5f;
81                 max_y = abs(y) + 0.5f;
82         }
83         assert(min_x >= 0.0f && max_x >= 0.0f);
84         assert(min_y >= 0.0f && max_y >= 0.0f);
85
86         float cell_height = max_y - min_y;
87         float cell_width = max_x - min_x;
88
89         if (min_x * min_x + min_y * min_y > r * r) {
90                 // Lower-left corner is outside the circle, so the entire cell is.
91                 return 0.0f;
92         }
93         if (max_x * max_x + max_y * max_y < r * r) {
94                 // Upper-right corner is inside the circle, so the entire cell is.
95                 return 1.0f;
96         }
97
98         // OK, so now we know the cell is partially covered by the circle:
99         //
100         //      \           .
101         //  -------------
102         // |####|#\      |
103         // |####|##|     |
104         //  -------------
105         //   A   ###|
106         //       ###|
107         //
108         // The edge of the circle is defined by x² + y² = r², 
109         // or x = sqrt(r² - y²) (since x is nonnegative).
110         // Find out where the curve crosses our given y values.
111         float mid_x1 = (max_y >= r) ? min_x : sqrt(r * r - max_y * max_y);
112         float mid_x2 = sqrt(r * r - min_y * min_y);
113         if (mid_x1 < min_x) {
114                 mid_x1 = min_x;
115         }
116         if (mid_x2 > max_x) {
117                 mid_x2 = max_x;
118         }
119         assert(mid_x1 >= min_x);
120         assert(mid_x2 >= mid_x1);
121         assert(max_x >= mid_x2);
122
123         // The area marked A in the figure above.
124         float covered_area = cell_height * (mid_x1 - min_x);
125
126         // The area marked B in the figure above. Note that the integral gives the entire
127         // shaded space down to zero, so we need to subtract the rectangle that does not
128         // belong to our cell.
129         covered_area += circle_integral(mid_x2, r) - circle_integral(mid_x1, r);
130         covered_area -= min_y * (mid_x2 - mid_x1);
131
132         assert(covered_area <= cell_width * cell_height);
133         return covered_area / (cell_width * cell_height);
134 }
135
136 // Compute a ⊙ b. Note that we compute the “full” convolution,
137 // ie., our matrix will be big enough to hold every nonzero element of the result.
138 MatrixXf convolve(const MatrixXf &a, const MatrixXf &b)
139 {
140         MatrixXf result(a.rows() + b.rows() - 1, a.cols() + b.cols() - 1);
141         for (int yr = 0; yr < result.rows(); ++yr) {
142                 for (int xr = 0; xr < result.cols(); ++xr) {
143                         float sum = 0.0f;
144
145                         // Given that x_b = x_r - x_a, find the values of x_a where
146                         // x_a is in [0, a_cols> and x_b is in [0, b_cols>. (y is similar.)
147                         //
148                         // The second demand gives:
149                         //
150                         //   0 <= x_r - x_a < b_cols
151                         //   0 >= x_a - x_r > -b_cols
152                         //   x_r >= x_a > x_r - b_cols
153                         int ya_min = yr - b.rows() + 1;
154                         int ya_max = yr;
155                         int xa_min = xr - b.rows() + 1;
156                         int xa_max = xr;
157
158                         // Now fit to the first demand.
159                         ya_min = std::max<int>(ya_min, 0);
160                         ya_max = std::min<int>(ya_max, a.rows() - 1);
161                         xa_min = std::max<int>(xa_min, 0);
162                         xa_max = std::min<int>(xa_max, a.cols() - 1);
163
164                         assert(ya_max >= ya_min);
165                         assert(xa_max >= xa_min);
166
167                         for (int ya = ya_min; ya <= ya_max; ++ya) {
168                                 for (int xa = xa_min; xa <= xa_max; ++xa) {
169                                         sum += a(ya, xa) * b(yr - ya, xr - xa);
170                                 }
171                         }
172
173                         result(yr, xr) = sum;
174                 }
175         }
176         return result;
177 }
178
179 // Similar to convolve(), but instead of assuming every element outside
180 // of b is zero, we make no such assumption and instead return only the
181 // elements where we know the right answer. (This is the only difference
182 // between the two.)
183 // This is the same as conv2(a, b, 'valid') in Octave.
184 //
185 // a must be the larger matrix of the two.
186 MatrixXf central_convolve(const MatrixXf &a, const MatrixXf &b)
187 {
188         assert(a.rows() >= b.rows());
189         assert(a.cols() >= b.cols());
190         MatrixXf result(a.rows() - b.rows() + 1, a.cols() - b.cols() + 1);
191         for (int yr = b.rows() - 1; yr < result.rows() + b.rows() - 1; ++yr) {
192                 for (int xr = b.cols() - 1; xr < result.cols() + b.cols() - 1; ++xr) {
193                         float sum = 0.0f;
194
195                         // Given that x_b = x_r - x_a, find the values of x_a where
196                         // x_a is in [0, a_cols> and x_b is in [0, b_cols>. (y is similar.)
197                         //
198                         // The second demand gives:
199                         //
200                         //   0 <= x_r - x_a < b_cols
201                         //   0 >= x_a - x_r > -b_cols
202                         //   x_r >= x_a > x_r - b_cols
203                         int ya_min = yr - b.rows() + 1;
204                         int ya_max = yr;
205                         int xa_min = xr - b.rows() + 1;
206                         int xa_max = xr;
207
208                         // Now fit to the first demand.
209                         ya_min = std::max<int>(ya_min, 0);
210                         ya_max = std::min<int>(ya_max, a.rows() - 1);
211                         xa_min = std::max<int>(xa_min, 0);
212                         xa_max = std::min<int>(xa_max, a.cols() - 1);
213
214                         assert(ya_max >= ya_min);
215                         assert(xa_max >= xa_min);
216
217                         for (int ya = ya_min; ya <= ya_max; ++ya) {
218                                 for (int xa = xa_min; xa <= xa_max; ++xa) {
219                                         sum += a(ya, xa) * b(yr - ya, xr - xa);
220                                 }
221                         }
222
223                         result(yr - b.rows() + 1, xr - b.cols() + 1) = sum;
224                 }
225         }
226         return result;
227 }
228
229 void print_matrix(const MatrixXf &m)
230 {
231         for (int y = 0; y < m.rows(); ++y) {
232                 for (int x = 0; x < m.cols(); ++x) {
233                         printf("%7.4f ", m(x, y));
234                 }
235                 printf("\n");
236         }
237 }
238
239 }  // namespace
240
241 void DeconvolutionSharpenEffect::set_gl_state(GLuint glsl_program_num, const std::string &prefix, unsigned *sampler_num)
242 {
243         Effect::set_gl_state(glsl_program_num, prefix, sampler_num);
244
245         assert(R >= 1);
246         assert(R <= 25);  // Same limit as Refocus.
247
248         printf("circular blur radius: %5.3f\n", circle_radius);
249         printf("gaussian blur radius: %5.3f\n", gaussian_radius);
250         printf("correlation:          %5.3f\n", correlation);
251         printf("noise factor:         %5.3f\n", noise);
252         printf("\n");
253
254         // Figure out the impulse response for the circular part of the blur.
255         MatrixXf circ_h(2 * R + 1, 2 * R + 1);
256         for (int y = -R; y <= R; ++y) { 
257                 for (int x = -R; x <= R; ++x) {
258                         circ_h(y + R, x + R) = circle_impulse_response(x, y, circle_radius);
259                 }
260         }
261
262         // Same, for the Gaussian part of the blur. We make this a lot larger
263         // since we're going to convolve with it soon, and it has infinite support
264         // (see comments for central_convolve()).
265         MatrixXf gaussian_h(4 * R + 1, 4 * R + 1);
266         for (int y = -2 * R; y <= 2 * R; ++y) { 
267                 for (int x = -2 * R; x <= 2 * R; ++x) {
268                         float val;
269                         if (gaussian_radius < 1e-3) {
270                                 val = (x == 0 && y == 0) ? 1.0f : 0.0f;
271                         } else {
272                                 float z = hypot(x, y) / gaussian_radius;
273                                 val = exp(-z * z);
274                         }
275                         gaussian_h(y + 2 * R, x + 2 * R) = val;
276                 }
277         }
278
279         // h, the (assumed) impulse response that we're trying to invert.
280         MatrixXf h = central_convolve(gaussian_h, circ_h);
281         assert(h.rows() == 2 * R + 1);
282         assert(h.cols() == 2 * R + 1);
283
284         // Normalize the impulse response.
285         float sum = 0.0f;
286         for (int y = 0; y < 2 * R + 1; ++y) {
287                 for (int x = 0; x < 2 * R + 1; ++x) {
288                         sum += h(y, x);
289                 }
290         }
291         for (int y = 0; y < 2 * R + 1; ++y) {
292                 for (int x = 0; x < 2 * R + 1; ++x) {
293                         h(y, x) /= sum;
294                 }
295         }
296
297         // r_uu, the (estimated/assumed) autocorrelation of the input signal (u).
298         // The signal is modelled a standard autoregressive process with the
299         // given correlation coefficient.
300         //
301         // We have to take a bit of care with the size of this matrix.
302         // The pow() function naturally has an infinite support (except for the
303         // degenerate case of correlation=0), but we have to chop it off
304         // somewhere. Since we convolve it with a 4*R+1 large matrix below,
305         // we need to make it twice as big as that, so that we have enough
306         // data to make r_vv valid. (central_convolve() effectively enforces
307         // that we get at least the right size.)
308         MatrixXf r_uu(8 * R + 1, 8 * R + 1);
309         for (int y = -4 * R; y <= 4 * R; ++y) { 
310                 for (int x = -4 * R; x <= 4 * R; ++x) {
311                         r_uu(x + 4 * R, y + 4 * R) = pow(correlation, hypot(x, y));
312                 }
313         }
314
315         // Estimate r_vv, the autocorrelation of the output signal v.
316         // Since we know that v = h ⊙ u and both are symmetrical,
317         // convolution and correlation are the same, and
318         // r_vv = v ⊙ v = (h ⊙ u) ⊙ (h ⊙ u) = (h ⊙ h) ⊙ r_uu.
319         MatrixXf r_vv = central_convolve(r_uu, convolve(h, h));
320         assert(r_vv.rows() == 4 * R + 1);
321         assert(r_vv.cols() == 4 * R + 1);
322
323         // Similarly, r_uv = u ⊙ v = u ⊙ (h ⊙ u) = h ⊙ r_uu.
324         //MatrixXf r_uv = central_convolve(r_uu, h).block(2 * R, 2 * R, 2 * R + 1, 2 * R + 1);
325         MatrixXf r_uu_center = r_uu.block(2 * R, 2 * R, 4 * R + 1, 4 * R + 1);
326         MatrixXf r_uv = central_convolve(r_uu_center, h);
327         assert(r_uv.rows() == 2 * R + 1);
328         assert(r_uv.cols() == 2 * R + 1);
329         
330         // Add the noise term (we assume the noise is uncorrelated,
331         // so it only affects the central element).
332         r_vv(2 * R, 2 * R) += noise;
333
334         // Now solve the Wiener-Hopf equations to find the deconvolution kernel g.
335         // Most texts show this only for the simpler 1D case:
336         //
337         // [ r_vv(0)  r_vv(1) r_vv(2) ... ] [ g(0) ]   [ r_uv(0) ]
338         // [ r_vv(-1) r_vv(0) ...         ] [ g(1) ] = [ r_uv(1) ]
339         // [ r_vv(-2) ...                 ] [ g(2) ]   [ r_uv(2) ]
340         // [ ...                          ] [ g(3) ]   [ r_uv(3) ]
341         //
342         // (Since r_vv is symmetrical, we can drop the minus signs.)
343         //
344         // Generally, row i of the matrix contains (dropping _vv for brevity):
345         //
346         // [ r(0-i) r(1-i) r(2-i) ... ]
347         //
348         // However, we have the 2D case. We flatten the vectors out to
349         // 1D quantities; this means we must think of the row number
350         // as a pair instead of as a scalar. Row (i,j) then contains:
351         //
352         // [ r(0-i,0-j) r(1-i,0-j) r(2-i,0-j) ... r(0-i,1-j) r_(1-i,1-j) r(2-i,1-j) ... ]
353         //
354         // g and r_uv are flattened in the same fashion.
355         //
356         // Note that even though this matrix is block Toeplitz, it is _not_ Toeplitz,
357         // and thus can not be inverted through the standard Levinson-Durbin method.
358         // There exists a block Levinson-Durbin method, which we may or may not
359         // want to use later. (Eigen's solvers are fast enough that for big matrices,
360         // the convolution operation and not the matrix solving is the bottleneck.)
361         //
362         // One thing we definitely want to use, though, is the symmetry properties.
363         // Since we know that g(i, j) = g(|i|, |j|), we can reduce the amount of
364         // unknowns to about 1/4th of the total size. The method is quite simple,
365         // as can be seen from the following toy equation system:
366         //
367         //   A x0 + B x1 + C x2 = y0
368         //   D x0 + E x1 + F x2 = y1
369         //   G x0 + H x1 + I x2 = y2
370         //
371         // If we now know that e.g. x0=x1 and y0=y1, we can rewrite this to
372         //
373         //   (A+B+D+E) x0 + (C+F) x2 = 2 y0
374         //   (G+H)     x0 + I x2     = y2
375         //
376         // This both increases accuracy and provides us with a very nice speed
377         // boost. We could have gone even further and went for 8-way symmetry
378         // like the shader does, but this is good enough right now.
379         MatrixXf M(MatrixXf::Zero((R + 1) * (R + 1), (R + 1) * (R + 1)));
380         MatrixXf r_uv_flattened(MatrixXf::Zero((R + 1) * (R + 1), 1));
381         for (int outer_i = 0; outer_i < 2 * R + 1; ++outer_i) {
382                 int folded_outer_i = abs(outer_i - R);
383                 for (int outer_j = 0; outer_j < 2 * R + 1; ++outer_j) {
384                         int folded_outer_j = abs(outer_j - R);
385                         int row = folded_outer_i * (R + 1) + folded_outer_j;
386                         for (int inner_i = 0; inner_i < 2 * R + 1; ++inner_i) {
387                                 int folded_inner_i = abs(inner_i - R);
388                                 for (int inner_j = 0; inner_j < 2 * R + 1; ++inner_j) {
389                                         int folded_inner_j = abs(inner_j - R);
390                                         int col = folded_inner_i * (R + 1) + folded_inner_j;
391                                         M(row, col) += r_vv((inner_i - R) - (outer_i - R) + 2 * R,
392                                                             (inner_j - R) - (outer_j - R) + 2 * R);
393                                 }
394                         }
395                         r_uv_flattened(row) += r_uv(outer_i, outer_j);
396                 }
397         }
398
399         LLT<MatrixXf> llt(M);
400         MatrixXf g_flattened = llt.solve(r_uv_flattened);
401         assert(g_flattened.rows() == (R + 1) * (R + 1)),
402         assert(g_flattened.cols() == 1);
403
404         // Normalize and de-flatten the deconvolution matrix.
405         MatrixXf g(R + 1, R + 1);
406         sum = 0.0f;
407         for (int i = 0; i < g_flattened.rows(); ++i) {
408                 int y = i / (R + 1);
409                 int x = i % (R + 1);
410                 if (y == 0 && x == 0) {
411                         sum += g_flattened(i);
412                 } else if (y == 0 || x == 0) {
413                         sum += 2.0f * g_flattened(i);
414                 } else {
415                         sum += 4.0f * g_flattened(i);
416                 }
417         }
418         for (int i = 0; i < g_flattened.rows(); ++i) {
419                 int y = i / (R + 1);
420                 int x = i % (R + 1);
421                 g(y, x) = g_flattened(i) / sum;
422         }
423
424         // Now encode it as uniforms, and pass it on to the shader.
425         // (Actually the shader only uses about half of the elements.)
426         float samples[4 * (R + 1) * (R + 1)];
427         for (int y = 0; y <= R; ++y) {
428                 for (int x = 0; x <= R; ++x) {
429                         int i = y * (R + 1) + x;
430                         samples[i * 4 + 0] = x / float(width);
431                         samples[i * 4 + 1] = y / float(height);
432                         samples[i * 4 + 2] = g(y, x);
433                         samples[i * 4 + 3] = 0.0f;
434                 }
435         }
436
437         set_uniform_vec4_array(glsl_program_num, prefix, "samples", samples, R * R);
438 }