src/fft_fftw.cpp

   1 #include "fft_fftw.h"
   2
   3 #ifdef OPENMP
   4 #include <omp.h>
   5 #endif
   6
   7 Fftw::Fftw(){}
   8
   9 fftwf_plan Fftw::create_plan_fwd(uint howmany) const
  10 {
  11     cv::Mat mat_in = cv::Mat::zeros(howmany * m_height, m_width, CV_32F);
  12     ComplexMat mat_out(m_height, m_width / 2 + 1, howmany);
  13     float *in = reinterpret_cast<float *>(mat_in.data);
  14     fftwf_complex *out = reinterpret_cast<fftwf_complex *>(mat_out.get_p_data());
  15
  16     int rank = 2;
  17     int n[] = {(int)m_height, (int)m_width};
  18     int idist = m_height * m_width, odist = m_height * (m_width / 2 + 1);
  19     int istride = 1, ostride = 1;
  20     int *inembed = NULL, *onembed = NULL;
  21
  22     return fftwf_plan_many_dft_r2c(rank, n, howmany, in, inembed, istride, idist, out, onembed, ostride, odist, FFTW_PATIENT);
  23 }
  24
  25 fftwf_plan Fftw::create_plan_inv(uint howmany) const
  26 {
  27     ComplexMat mat_in(m_height, m_width / 2 + 1, howmany);
  28     cv::Mat mat_out = cv::Mat::zeros(howmany * m_height, m_width, CV_32F);
  29     fftwf_complex *in = reinterpret_cast<fftwf_complex *>(mat_in.get_p_data());
  30     float *out = reinterpret_cast<float *>(mat_out.data);
  31
  32     int rank = 2;
  33     int n[] = {(int)m_height, (int)m_width};
  34     int idist = m_height * (m_width / 2 + 1), odist = m_height * m_width;
  35     int istride = 1, ostride = 1;
  36     int *inembed = nullptr, *onembed = nullptr;
  37
  38     return fftwf_plan_many_dft_c2r(rank, n, howmany, in, inembed, istride, idist, out, onembed, ostride, odist, FFTW_PATIENT);
  39 }
  40
  41 void Fftw::init(unsigned width, unsigned height, unsigned num_of_feats, unsigned num_of_scales)
  42 {
  43     Fft::init(width, height, num_of_feats, num_of_scales);
  44
  45 #if !defined(CUFFTW)
  46     fftw_init_threads();
  47   #if defined(OPENMP)
  48     fftw_plan_with_nthreads(omp_get_max_threads());
  49   #else
  50     fftw_plan_with_nthreads(4);
  51   #endif
  52 #endif
  53
  54 #ifndef CUFFTW
  55     std::cout << "FFT: FFTW" << std::endl;
  56 #else
  57     std::cout << "FFT: cuFFTW" << std::endl;
  58 #endif
  59     fftwf_cleanup();
  60
  61     plan_f = create_plan_fwd(1);
  62     plan_fw = create_plan_fwd(m_num_of_feats);
  63     plan_i_1ch = create_plan_inv(1);
  64
  65 #ifdef BIG_BATCH
  66     plan_f_all_scales = create_plan_fwd(m_num_of_scales);
  67     plan_fw_all_scales = create_plan_fwd(m_num_of_scales * m_num_of_feats);
  68     plan_i_all_scales = create_plan_inv(m_num_of_scales);
  69 #endif
  70 }
  71
  72 void Fftw::set_window(const MatDynMem &window)
  73 {
  74     Fft::set_window(window);
  75     m_window = window;
  76 }
  77
  78 void Fftw::forward(const MatScales &real_input, ComplexMat &complex_result)
  79 {
  80     Fft::forward(real_input, complex_result);
  81
  82     if (real_input.size[0] == 1)
  83         fftwf_execute_dft_r2c(plan_f, reinterpret_cast<float *>(real_input.data),
  84                               reinterpret_cast<fftwf_complex *>(complex_result.get_p_data()));
  85 #ifdef BIG_BATCH
  86     else
  87         fftwf_execute_dft_r2c(plan_f_all_scales, reinterpret_cast<float *>(real_input.data),
  88                               reinterpret_cast<fftwf_complex *>(complex_result.get_p_data()));
  89 #endif
  90 }
  91
  92 void Fftw::forward_window(MatScaleFeats  &feat, ComplexMat & complex_result, MatScaleFeats &temp)
  93 {
  94     Fft::forward_window(feat, complex_result, temp);
  95
  96     uint n_scales = feat.size[0];
  97     for (uint s = 0; s < n_scales; ++s) {
  98         for (uint ch = 0; ch < uint(feat.size[1]); ++ch) {
  99             cv::Mat feat_plane = feat.plane(s, ch);
 100             cv::Mat temp_plane = temp.plane(s, ch);
 101             temp_plane = feat_plane.mul(m_window);
 102         }
 103     }
 104
 105     float *in = temp.ptr<float>();
 106     fftwf_complex *out = reinterpret_cast<fftwf_complex *>(complex_result.get_p_data());
 107
 108     if (n_scales == 1)
 109         fftwf_execute_dft_r2c(plan_fw, in, out);
 110 #ifdef BIG_BATCH
 111     else
 112         fftwf_execute_dft_r2c(plan_fw_all_scales, in, out);
 113 #endif
 114 }
 115
 116 void Fftw::inverse(ComplexMat &complex_input, MatScales &real_result)
 117 {
 118     Fft::inverse(complex_input, real_result);
 119
 120     int n_channels = complex_input.n_channels;
 121     fftwf_complex *in = reinterpret_cast<fftwf_complex *>(complex_input.get_p_data());
 122     float *out = real_result.ptr<float>();
 123
 124     if (n_channels == 1)
 125         fftwf_execute_dft_c2r(plan_i_1ch, in, out);
 126 #ifdef BIG_BATCH
 127     else
 128         fftwf_execute_dft_c2r(plan_i_all_scales, in, out);
 129 #endif
 130     real_result *= 1.0 / (m_width * m_height);
 131 }
 132
 133 Fftw::~Fftw()
 134 {
 135     fftwf_destroy_plan(plan_f);
 136     fftwf_destroy_plan(plan_fw);
 137     fftwf_destroy_plan(plan_i_1ch);
 138
 139 #ifdef BIG_BATCH
 140     fftwf_destroy_plan(plan_f_all_scales);
 141     fftwf_destroy_plan(plan_fw_all_scales);
 142     fftwf_destroy_plan(plan_i_all_scales);
 143 #endif
 144 }