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