src/fft_cufft.h

   1 #ifndef FFT_CUDA_H
   2 #define FFT_CUDA_H
   3
   4
   5 #include <cufft.h>
   6 #include <cuda_runtime.h>
   7
   8 #include "fft.h"
   9 #include "cuda/cuda_error_check.cuh"
  10 #include "pragmas.h"
  11
  12 #if CV_MAJOR_VERSION == 2
  13   #include <opencv2/gpu/gpu.hpp>
  14   #define CUDA cv::gpu
  15 #else
  16   #include "opencv2/opencv.hpp"
  17   #define CUDA cv::cuda
  18 #endif
  19
  20 struct ThreadCtx;
  21
  22 class cuFFT : public Fft
  23 {
  24 public:
  25     void init(unsigned width, unsigned height, unsigned num_of_feats, unsigned num_of_scales, bool big_batch_mode) override;
  26     void set_window(const cv::Mat & window) override;
  27     void forward(const cv::Mat & real_input, ComplexMat & complex_result, float *real_input_arr, cudaStream_t  stream) override;
  28     void forward_window(std::vector<cv::Mat> patch_feats, ComplexMat & complex_result, cv::Mat & fw_all, float *real_input_arr, cudaStream_t stream) override;
  29     void inverse(ComplexMat &  complex_input, cv::Mat & real_result, float *real_result_arr, cudaStream_t stream) override;
  30     ~cuFFT() override;
  31 private:
  32     cv::Mat m_window;
  33     unsigned m_width, m_height, m_num_of_feats, m_num_of_scales;
  34     bool m_big_batch_mode;
  35     cufftHandle plan_f, plan_f_all_scales, plan_fw, plan_fw_all_scales, plan_i_features,
  36      plan_i_features_all_scales, plan_i_1ch, plan_i_1ch_all_scales;
  37 };
  38
  39 #endif // FFT_CUDA_H