src/fft_cufft.h

   1
   2 #ifndef FFT_CUDA_H
   3 #define FFT_CUDA_H
   4
   5 #include "fft.h"
   6
   7 #if CV_MAJOR_VERSION == 2
   8   #include <opencv2/gpu/gpu.hpp>
   9   #define CUDA cv::gpu
  10 #else
  11   #include "opencv2/opencv.hpp"
  12   #define CUDA cv::cuda
  13 #endif
  14
  15 #include <cufft.h>
  16 #include <cuda_runtime.h>
  17
  18 class cuFFT : public Fft
  19 {
  20 public:
  21     cuFFT();
  22     void init(unsigned width, unsigned height, unsigned num_of_feats, unsigned num_of_scales, bool big_batch_mode) override;
  23     void set_window(const cv::Mat &window) override;
  24     ComplexMat forward(const cv::Mat &input) override;
  25     ComplexMat forward_window(const std::vector<cv::Mat> &input) override;
  26     cv::Mat inverse(const ComplexMat &inputf) override;
  27     ~cuFFT() override;
  28 private:
  29     cv::Mat m_window;
  30     unsigned m_width, m_height, m_num_of_feats, m_num_of_scales, m_num_of_streams;
  31     bool m_big_batch_mode;
  32     cudaStream_t streams[4];
  33     cufftHandle plan_f, plan_f_all_scales, plan_fw, plan_fw_all_scales, plan_i_features,
  34      plan_i_features_all_scales, plan_i_1ch, plan_i_1ch_all_scales;
  35     float *data_f, *data_f_all_scales, *data_fw, *data_fw_d, *data_fw_all_scales, *data_fw_all_scales_d, *data_i_features, *data_i_features_d,
  36           *data_i_features_all_scales, *data_i_features_all_scales_d, *data_i_1ch, *data_i_1ch_d, *data_i_1ch_all_scales, *data_i_1ch_all_scales_d;
  37 };
  38
  39 #endif // FFT_CUDA_H