src/fft_fftw.cpp

   1 #include "fft_fftw.h"
   2
   3 #include "fft.h"
   4
   5 #ifdef OPENMP
   6   #include <omp.h>
   7 #endif
   8
   9 #if !defined(ASYNC) && !defined(OPENMP) && !defined(CUFFTW)
  10 #define FFTW_PLAN_WITH_THREADS() fftw_plan_with_nthreads(m_num_threads);
  11 #elif defined(OPENMP) && !defined(CUFFTW)
  12 #define FFTW_PLAN_WITH_THREADS() fftw_plan_with_nthreads(omp_get_max_threads());
  13 #else
  14 #define FFTW_PLAN_WITH_THREADS()
  15 #endif
  16
  17 Fftw::Fftw()
  18     : m_num_threads(2)
  19 {
  20 }
  21
  22 Fftw::Fftw(int num_threads)
  23     : m_num_threads(num_threads)
  24 {
  25 }
  26
  27 void Fftw::init(unsigned width, unsigned height, unsigned num_of_feats, unsigned num_of_scales, bool big_batch_mode)
  28 {
  29     m_width = width;
  30     m_height = height;
  31     m_num_of_feats = num_of_feats;
  32     m_num_of_scales = num_of_scales;
  33     m_big_batch_mode = big_batch_mode;
  34
  35 #if (!defined(ASYNC) && !defined(CUFFTW))|| defined(OPENMP)
  36     fftw_init_threads();
  37 #endif //OPENMP
  38
  39 #ifndef CUFFTW
  40     std::cout << "FFT: FFTW" << std::endl;
  41 #else
  42     std::cout << "FFT: cuFFTW" << std::endl;
  43 #endif
  44     //FFT forward one scale
  45     {
  46         cv::Mat in_f = cv::Mat::zeros(m_height, m_width, CV_32FC1);
  47         ComplexMat out_f(m_height, m_width / 2 + 1, 1);
  48         plan_f = fftwf_plan_dft_r2c_2d(m_height, m_width,
  49                                        reinterpret_cast<float*>(in_f.data),
  50                                        reinterpret_cast<fftwf_complex*>(out_f.get_p_data()),
  51                                        FFTW_MEASURE|FFTW_PATIENT);
  52     }
  53     //FFT forward all scales
  54     if (m_num_of_scales > 1 && m_big_batch_mode) {
  55         cv::Mat in_f_all = cv::Mat::zeros(m_height*m_num_of_scales, m_width, CV_32F);
  56         ComplexMat out_f_all(m_height, m_width / 2 + 1, m_num_of_scales);
  57         float *in = reinterpret_cast<float*>(in_f_all.data);
  58         fftwf_complex *out = reinterpret_cast<fftwf_complex*>(out_f_all.get_p_data());
  59         int rank = 2;
  60         int n[] = {(int)m_height, (int)m_width};
  61         int howmany = m_num_of_scales;
  62         int idist = m_height*m_width, odist = m_height*(m_width/2+1);
  63         int istride = 1, ostride = 1;
  64         int *inembed = NULL, *onembed = NULL;
  65
  66         FFTW_PLAN_WITH_THREADS();
  67         plan_f_all_scales = fftwf_plan_many_dft_r2c(rank, n, howmany,
  68                                                     in, inembed, istride, idist,
  69                                                     out, onembed, ostride, odist,
  70                                                     FFTW_MEASURE|FFTW_PATIENT);
  71     }
  72     //FFT forward window one scale
  73     {
  74         cv::Mat in_fw = cv::Mat::zeros(m_height * m_num_of_feats, m_width, CV_32F);
  75         ComplexMat out_fw(m_height, m_width / 2 + 1, m_num_of_feats);
  76         float *in = reinterpret_cast<float*>(in_fw.data);
  77         fftwf_complex *out = reinterpret_cast<fftwf_complex*>(out_fw.get_p_data());
  78         int rank = 2;
  79         int n[] = {(int)m_height, (int)m_width};
  80         int howmany = m_num_of_feats;
  81         int idist = m_height*m_width, odist = m_height*(m_width/2+1);
  82         int istride = 1, ostride = 1;
  83         int *inembed = NULL, *onembed = NULL;
  84
  85         FFTW_PLAN_WITH_THREADS();
  86         plan_fw = fftwf_plan_many_dft_r2c(rank, n, howmany,
  87                                           in, inembed, istride, idist,
  88                                           out, onembed, ostride, odist,
  89                                           FFTW_MEASURE|FFTW_PATIENT);
  90     }
  91     //FFT forward window all scales all feats
  92     if (m_num_of_scales > 1 && m_big_batch_mode) {
  93         cv::Mat in_all = cv::Mat::zeros(m_height * (m_num_of_scales*m_num_of_feats), m_width, CV_32F);
  94         ComplexMat out_all(m_height, m_width / 2 + 1, m_num_of_scales*m_num_of_feats);
  95         float *in = reinterpret_cast<float*>(in_all.data);
  96         fftwf_complex *out = reinterpret_cast<fftwf_complex*>(out_all.get_p_data());
  97         int rank = 2;
  98         int n[] = {(int)m_height, (int)m_width};
  99         int howmany = m_num_of_scales*m_num_of_feats;
 100         int idist = m_height*m_width, odist = m_height*(m_width/2+1);
 101         int istride = 1, ostride = 1;
 102         int *inembed = NULL, *onembed = NULL;
 103
 104         FFTW_PLAN_WITH_THREADS();
 105         plan_fw_all_scales = fftwf_plan_many_dft_r2c(rank, n, howmany,
 106                                                      in,  inembed, istride, idist,
 107                                                      out, onembed, ostride, odist,
 108                                                      FFTW_MEASURE|FFTW_PATIENT);
 109     }
 110     //FFT inverse one scale
 111     {
 112         ComplexMat in_i(m_height,m_width,m_num_of_feats);
 113         cv::Mat out_i = cv::Mat::zeros(m_height, m_width, CV_32FC(m_num_of_feats));
 114         fftwf_complex *in = reinterpret_cast<fftwf_complex*>(in_i.get_p_data());
 115         float *out = reinterpret_cast<float*>(out_i.data);
 116         int rank = 2;
 117         int n[] = {(int)m_height, (int)m_width};
 118         int howmany = m_num_of_feats;
 119         int idist = m_height*(m_width/2+1), odist = 1;
 120         int istride = 1, ostride = m_num_of_feats;
 121         int inembed[] = {(int)m_height, (int)m_width/2+1}, *onembed = n;
 122
 123         FFTW_PLAN_WITH_THREADS();
 124         plan_i_features = fftwf_plan_many_dft_c2r(rank, n, howmany,
 125                                                   in,  inembed, istride, idist,
 126                                                   out, onembed, ostride, odist,
 127                                                   FFTW_MEASURE|FFTW_PATIENT);
 128     }
 129     //FFT inverse all scales
 130     if (m_num_of_scales > 1 && m_big_batch_mode) {
 131         ComplexMat in_i_all(m_height,m_width,m_num_of_feats*m_num_of_scales);
 132         cv::Mat out_i_all = cv::Mat::zeros(m_height, m_width, CV_32FC(m_num_of_feats*m_num_of_scales));
 133         fftwf_complex *in = reinterpret_cast<fftwf_complex*>(in_i_all.get_p_data());
 134         float *out = reinterpret_cast<float*>(out_i_all.data);
 135         int rank = 2;
 136         int n[] = {(int)m_height, (int)m_width};
 137         int howmany = m_num_of_feats*m_num_of_scales;
 138         int idist = m_height*(m_width/2+1), odist = 1;
 139         int istride = 1, ostride = m_num_of_feats*m_num_of_scales;
 140         int inembed[] = {(int)m_height, (int)m_width/2+1}, *onembed = n;
 141
 142         FFTW_PLAN_WITH_THREADS();
 143         plan_i_features_all_scales = fftwf_plan_many_dft_c2r(rank, n, howmany,
 144                                                              in,  inembed, istride, idist,
 145                                                              out, onembed, ostride, odist,
 146                                                              FFTW_MEASURE|FFTW_PATIENT);
 147     }
 148     //FFT inver one channel one scale
 149     if(!m_big_batch_mode) {
 150         ComplexMat in_i1(m_height,m_width,1);
 151         cv::Mat out_i1 = cv::Mat::zeros(m_height, m_width, CV_32FC1);
 152         fftwf_complex *in = reinterpret_cast<fftwf_complex*>(in_i1.get_p_data());
 153         float *out = reinterpret_cast<float*>(out_i1.data);
 154         int rank = 2;
 155         int n[] = {(int)m_height, (int)m_width};
 156         int howmany = 1;
 157         int idist = m_height*(m_width/2+1), odist = 1;
 158         int istride = 1, ostride = 1;
 159         int inembed[] = {(int)m_height, (int)m_width/2+1}, *onembed = n;
 160
 161         FFTW_PLAN_WITH_THREADS();
 162         plan_i_1ch = fftwf_plan_many_dft_c2r(rank, n, howmany,
 163                                              in,  inembed, istride, idist,
 164                                              out, onembed, ostride, odist,
 165                                              FFTW_MEASURE|FFTW_PATIENT);
 166     }
 167     //FFT inver one channel all scales
 168     if (m_num_of_scales > 1 && m_big_batch_mode) {
 169         ComplexMat in_i1_all(m_height,m_width,m_num_of_scales);
 170         cv::Mat out_i1_all = cv::Mat::zeros(m_height, m_width, CV_32FC(m_num_of_scales));
 171         fftwf_complex *in = reinterpret_cast<fftwf_complex*>(in_i1_all.get_p_data());
 172         float *out = reinterpret_cast<float*>(out_i1_all.data);
 173         int rank = 2;
 174         int n[] = {(int)m_height, (int)m_width};
 175         int howmany = m_num_of_scales;
 176         int idist = m_height*(m_width/2+1), odist = 1;
 177         int istride = 1, ostride = m_num_of_scales;
 178         int inembed[] = {(int)m_height, (int)m_width/2+1}, *onembed = n;
 179
 180         FFTW_PLAN_WITH_THREADS();
 181         plan_i_1ch_all_scales = fftwf_plan_many_dft_c2r(rank, n, howmany,
 182                                                         in,  inembed, istride, idist,
 183                                                         out, onembed, ostride, odist,
 184                                                         FFTW_MEASURE|FFTW_PATIENT);
 185     }
 186 }
 187
 188 void Fftw::set_window(const cv::Mat &window)
 189 {
 190     m_window = window;
 191 }
 192
 193 ComplexMat Fftw::forward(const cv::Mat &input)
 194 {
 195     ComplexMat complex_result;
 196     if(m_big_batch_mode && input.rows == (int)(m_height*m_num_of_scales)){
 197         complex_result.create(m_height, m_width / 2 + 1, m_num_of_scales);
 198         fftwf_execute_dft_r2c(plan_f_all_scales, reinterpret_cast<float*>(input.data),
 199                               reinterpret_cast<fftwf_complex*>(complex_result.get_p_data()));
 200     } else {
 201         complex_result.create(m_height, m_width / 2 + 1, 1);
 202         fftwf_execute_dft_r2c(plan_f, reinterpret_cast<float*>(input.data),
 203                               reinterpret_cast<fftwf_complex*>(complex_result.get_p_data()));
 204     }
 205     return complex_result;
 206 }
 207
 208 ComplexMat Fftw::forward_raw(float *input, bool all_scales)
 209 {
 210     ComplexMat dummy;
 211     return dummy;
 212 }
 213
 214 ComplexMat Fftw::forward_window(const std::vector<cv::Mat> &input)
 215 {
 216     int n_channels = input.size();
 217     cv::Mat in_all(m_height * n_channels, m_width, CV_32F);
 218     for (int i = 0; i < n_channels; ++i) {
 219         cv::Mat in_roi(in_all, cv::Rect(0, i*m_height, m_width, m_height));
 220         in_roi = input[i].mul(m_window);
 221     }
 222     ComplexMat result;
 223     if(n_channels > (int) m_num_of_feats)
 224         result.create(m_height, m_width/2 + 1, n_channels,m_num_of_scales);
 225     else
 226         result.create(m_height, m_width/2 + 1, n_channels);
 227
 228     float *in = reinterpret_cast<float*>(in_all.data);
 229     fftwf_complex *out = reinterpret_cast<fftwf_complex*>(result.get_p_data());
 230
 231     if (n_channels <= (int) m_num_of_feats)
 232         fftwf_execute_dft_r2c(plan_fw, in, out);
 233     else
 234         fftwf_execute_dft_r2c(plan_fw_all_scales, in, out);
 235
 236     return result;
 237 }
 238
 239 cv::Mat Fftw::inverse(const ComplexMat &input)
 240 {
 241     int n_channels = input.n_channels;
 242     cv::Mat real_result(m_height, m_width, CV_32FC(n_channels));
 243     fftwf_complex *in = reinterpret_cast<fftwf_complex*>(input.get_p_data());
 244     float *out = reinterpret_cast<float*>(real_result.data);
 245
 246     if(n_channels == 1)
 247         fftwf_execute_dft_c2r(plan_i_1ch, in, out);
 248     else if(m_big_batch_mode && n_channels == (int) m_num_of_scales)
 249         fftwf_execute_dft_c2r(plan_i_1ch_all_scales, in, out);
 250     else if(m_big_batch_mode && n_channels == (int) m_num_of_feats * (int) m_num_of_scales)
 251         fftwf_execute_dft_c2r(plan_i_features_all_scales, in, out);
 252     else
 253         fftwf_execute_dft_c2r(plan_i_features, in, out);
 254
 255     return real_result/(m_width*m_height);
 256 }
 257
 258 float* Fftw::inverse_raw(const ComplexMat &input)
 259 {
 260     return nullptr;
 261 }
 262
 263 Fftw::~Fftw()
 264 {
 265     fftwf_destroy_plan(plan_f);
 266     fftwf_destroy_plan(plan_fw);
 267     fftwf_destroy_plan(plan_i_features);
 268     fftwf_destroy_plan(plan_i_1ch);
 269
 270     if (m_big_batch_mode) {
 271         fftwf_destroy_plan(plan_f_all_scales);
 272         fftwf_destroy_plan(plan_i_features_all_scales);
 273         fftwf_destroy_plan(plan_fw_all_scales);
 274         fftwf_destroy_plan(plan_i_1ch_all_scales);
 275     }
 276 }