src/fft_fftw.cpp

   1 #include "fft_fftw.h"
   2
   3 #include "fft.h"
   4 #ifndef CUFFTW
   5   #include <fftw3.h>
   6 #else
   7   #include <cufftw.h>
   8 #endif //CUFFTW
   9
  10 #if defined(OPENMP)
  11   #include <omp.h>
  12 #endif
  13
  14 Fftw::Fftw()
  15 {
  16 }
  17
  18 void Fftw::init(unsigned width, unsigned height)
  19 {
  20     m_width = width;
  21     m_height = height;
  22
  23 #if defined(OPENMP)
  24     fftw_init_threads();
  25     fftw_plan_with_nthreads(omp_get_max_threads());
  26 #endif //OPENMP
  27
  28 #ifndef CUFFTW
  29     std::cout << "FFT: FFTW" << std::endl;
  30 #else
  31     std::cout << "FFT: cuFFTW" << std::endl;
  32 #endif
  33 }
  34
  35 void Fftw::set_window(const cv::Mat &window)
  36 {
  37     m_window = window;
  38 }
  39
  40 ComplexMat Fftw::forward(const cv::Mat &input)
  41 {
  42     cv::Mat complex_result(m_height, m_width / 2 + 1, CV_32FC2);
  43     fftwf_plan plan = fftwf_plan_dft_r2c_2d(m_height, m_width,
  44                                             reinterpret_cast<float*>(input.data),
  45                                             reinterpret_cast<fftwf_complex*>(complex_result.data),
  46                                             FFTW_ESTIMATE);
  47     fftwf_execute(plan);
  48     fftwf_destroy_plan(plan);
  49     return ComplexMat(complex_result);
  50 }
  51
  52 ComplexMat Fftw::forward_window(const std::vector<cv::Mat> &input)
  53 {
  54     int n_channels = input.size();
  55     cv::Mat in_all(m_height * n_channels, m_width, CV_32F);
  56     for (int i = 0; i < n_channels; ++i) {
  57         cv::Mat in_roi(in_all, cv::Rect(0, i*m_height, m_width, m_height));
  58         in_roi = input[i].mul(m_window);
  59     }
  60     cv::Mat complex_result(n_channels*m_height, m_width/2+1, CV_32FC2);
  61
  62     int rank = 2;
  63     int n[] = {(int)m_height, (int)m_width};
  64     int howmany = n_channels;
  65     int idist = m_height*m_width, odist = m_height*(m_width/2+1);
  66     int istride = 1, ostride = 1;
  67     int *inembed = NULL, *onembed = NULL;
  68     float *in = reinterpret_cast<float*>(in_all.data);
  69     fftwf_complex *out = reinterpret_cast<fftwf_complex*>(complex_result.data);
  70
  71     fftwf_plan plan = fftwf_plan_many_dft_r2c(rank, n, howmany,
  72                                               in,  inembed, istride, idist,
  73                                               out, onembed, ostride, odist,
  74                                               FFTW_ESTIMATE);
  75     fftwf_execute(plan);
  76     fftwf_destroy_plan(plan);
  77
  78     ComplexMat result(m_height, m_width/2 + 1, n_channels);
  79     for (int i = 0; i < n_channels; ++i)
  80         result.set_channel(i, complex_result(cv::Rect(0, i*m_height, m_width/2+1, m_height)));
  81
  82     return result;
  83 }
  84
  85 cv::Mat Fftw::inverse(const ComplexMat &inputf)
  86 {
  87     int n_channels = inputf.n_channels;
  88     cv::Mat real_result(m_height, m_width, CV_32FC(n_channels));
  89     cv::Mat complex_vconcat = inputf.to_vconcat_mat();
  90
  91     int rank = 2;
  92     int n[] = {(int)m_height, (int)m_width};
  93     int howmany = n_channels;
  94     int idist = m_height*(m_width/2+1), odist = 1;
  95     int istride = 1, ostride = n_channels;
  96 #ifndef CUFFTW
  97     int *inembed = NULL, *onembed = NULL;
  98 #else
  99     int inembed[2];
 100     int onembed[2];
 101     inembed[1] = m_width/2+1, onembed[1] = m_width;
 102 #endif
 103     fftwf_complex *in = reinterpret_cast<fftwf_complex*>(complex_vconcat.data);
 104     float *out = reinterpret_cast<float*>(real_result.data);
 105
 106     fftwf_plan plan = fftwf_plan_many_dft_c2r(rank, n, howmany,
 107                                               in,  inembed, istride, idist,
 108                                               out, onembed, ostride, odist,
 109                                               FFTW_ESTIMATE);
 110     fftwf_execute(plan);
 111     fftwf_destroy_plan(plan);
 112
 113     return real_result/(m_width*m_height);
 114 }
 115
 116 Fftw::~Fftw()
 117 {
 118 }