git.sesse.net Git - movit/blob - fft_input.cpp

   1 #include <string.h>
   2 #include <assert.h>
   3 #include <epoxy/gl.h>
   4 #include <fftw3.h>
   5
   6 #include "effect_util.h"
   7 #include "fp16.h"
   8 #include "fft_input.h"
   9 #include "resource_pool.h"
  10 #include "util.h"
  11
  12 using namespace std;
  13
  14 namespace movit {
  15
  16 FFTInput::FFTInput(unsigned width, unsigned height)
  17         : texture_num(0),
  18           fft_width(width),
  19           fft_height(height),
  20           convolve_width(width),
  21           convolve_height(height),
  22           pixel_data(NULL)
  23 {
  24         register_int("fft_width", &fft_width);
  25         register_int("fft_height", &fft_height);
  26         register_uniform_sampler2d("tex", &uniform_tex);
  27 }
  28
  29 FFTInput::~FFTInput()
  30 {
  31         if (texture_num != 0) {
  32                 resource_pool->release_2d_texture(texture_num);
  33         }
  34 }
  35
  36 void FFTInput::set_gl_state(GLuint glsl_program_num, const string& prefix, unsigned *sampler_num)
  37 {
  38         glActiveTexture(GL_TEXTURE0 + *sampler_num);
  39         check_error();
  40
  41         if (texture_num == 0) {
  42                 assert(pixel_data != NULL);
  43
  44                 // Do the FFT. Our FFTs should typically be small enough and
  45                 // the data changed often enough that FFTW_ESTIMATE should be
  46                 // quite OK. Otherwise, we'd need to worry about caching these
  47                 // plans (possibly including FFTW wisdom).
  48                 fftw_complex *in = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fft_width * fft_height);
  49                 fftw_complex *out = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * fft_width * fft_height);
  50                 fftw_plan p = fftw_plan_dft_2d(fft_height, fft_width, in, out, FFTW_FORWARD, FFTW_ESTIMATE);
  51
  52                 // Zero pad.
  53                 for (int i = 0; i < fft_height * fft_width; ++i) {
  54                         in[i][0] = 0.0;
  55                         in[i][1] = 0.0;
  56                 }
  57                 for (unsigned y = 0; y < convolve_height; ++y) {
  58                         for (unsigned x = 0; x < convolve_width; ++x) {
  59                                 int i = y * fft_width + x;
  60                                 in[i][0] = pixel_data[y * convolve_width + x];
  61                                 in[i][1] = 0.0;
  62                         }
  63                 }
  64
  65                 fftw_execute(p);
  66
  67                 // Convert to fp16.
  68                 fp16_int_t *kernel = new fp16_int_t[fft_width * fft_height * 2];
  69                 for (int i = 0; i < fft_width * fft_height; ++i) {
  70                         kernel[i * 2 + 0] = fp64_to_fp16(out[i][0]);
  71                         kernel[i * 2 + 1] = fp64_to_fp16(out[i][1]);
  72                 }
  73
  74                 // (Re-)upload the texture.
  75                 texture_num = resource_pool->create_2d_texture(GL_RG16F, fft_width, fft_height);
  76                 glBindTexture(GL_TEXTURE_2D, texture_num);
  77                 check_error();
  78                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST);
  79                 check_error();
  80                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_NEAREST);
  81                 check_error();
  82                 glPixelStorei(GL_UNPACK_ALIGNMENT, 1);
  83                 check_error();
  84                 glTexSubImage2D(GL_TEXTURE_2D, 0, 0, 0, fft_width, fft_height, GL_RG, GL_HALF_FLOAT, kernel);
  85                 check_error();
  86                 glPixelStorei(GL_UNPACK_ROW_LENGTH, 0);
  87                 check_error();
  88                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_REPEAT);
  89                 check_error();
  90                 glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_REPEAT);
  91                 check_error();
  92
  93                 fftw_free(in);
  94                 fftw_free(out);
  95                 delete[] kernel;
  96         } else {
  97                 glBindTexture(GL_TEXTURE_2D, texture_num);
  98                 check_error();
  99         }
 100
 101         // Bind it to a sampler.
 102         uniform_tex = *sampler_num;
 103         ++*sampler_num;
 104 }
 105
 106 string FFTInput::output_fragment_shader()
 107 {
 108         return string("#define FIXUP_SWAP_RB 0\n#define FIXUP_RED_TO_GRAYSCALE 0\n") +
 109                 read_file("flat_input.frag");
 110 }
 111
 112 void FFTInput::invalidate_pixel_data()
 113 {
 114         if (texture_num != 0) {
 115                 resource_pool->release_2d_texture(texture_num);
 116                 texture_num = 0;
 117         }
 118 }
 119
 120 bool FFTInput::set_int(const std::string& key, int value)
 121 {
 122         if (key == "needs_mipmaps") {
 123                 // We cannot supply mipmaps; it would not make any sense for FFT data.
 124                 return (value == 0);
 125         }
 126         if (key == "fft_width") {
 127                 if (value < int(convolve_width)) {
 128                         return false;
 129                 }
 130                 invalidate_pixel_data();
 131         }
 132         if (key == "fft_height") {
 133                 if (value < int(convolve_height)) {
 134                         return false;
 135                 }
 136                 invalidate_pixel_data();
 137         }
 138         return Effect::set_int(key, value);
 139 }
 140
 141 }  // namespace movit