src/fft_fftw.cpp

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