Unify CPU and GPU implementations of ComplexMat

[hercules2020/kcf.git] / src / threadctx.hpp

diff --git a/src/threadctx.hpp b/src/threadctx.hpp
index de9638e32ad4d32d9cebda4242a3842c8e6d54b5..982905b96699b5928d39eb2ba158fc6920c2120f 100644 (file)
--- a/src/threadctx.hpp
+++ b/src/threadctx.hpp
@@ -3,104 +3,64 @@
 
 #include <future>
 #include "dynmem.hpp"
-
-#ifdef CUFFT
-#include "complexmat.cuh"
-#else
+#include "kcf.h"
 #include "complexmat.hpp"
-#ifndef CUFFTW
-// For compatibility reasons between CuFFT and FFTW, OpenCVfft versions.
-typedef int *cudaStream_t;
-#endif
-#endif
+
+class KCF_Tracker;
 
 struct ThreadCtx {
   public:
-    ThreadCtx(cv::Size roi, uint num_of_feats, double scale, uint num_of_scales)
-        : scale(scale)
-    {
-        this->xf_sqr_norm = DynMem(num_of_scales * sizeof(float));
-        this->yf_sqr_norm = DynMem(sizeof(float));
-
-        uint cells_size = roi.width * roi.height * sizeof(float);
-
-#if  !defined(BIG_BATCH) && defined(CUFFT) && (defined(ASYNC) || defined(OPENMP))
-        CudaSafeCall(cudaStreamCreate(&this->stream));
+    ThreadCtx(cv::Size roi, uint num_features
+#ifdef BIG_BATCH
+              , uint num_scales
+#else
+              , double scale
 #endif
+             )
+        : roi(roi)
+        , num_features(num_features)
+        , num_scales(IF_BIG_BATCH(num_scales, 1))
+#ifndef BIG_BATCH
+        , scale(scale)
+#endif
+    {}
 
-#if defined(CUFFT) || defined(FFTW)
-        this->gauss_corr_res = DynMem(cells_size * num_of_scales);
-        this->data_features = DynMem(cells_size * num_of_feats);
+    ThreadCtx(ThreadCtx &&) = default;
 
-        uint width_freq = roi.width / 2 + 1;
+    void track(const KCF_Tracker &kcf, cv::Mat &input_rgb, cv::Mat &input_gray);
+private:
+    cv::Size roi;
+    uint num_features;
+    uint num_scales;
+    cv::Size freq_size = Fft::freq_size(roi);
 
-        this->in_all = cv::Mat(roi.height * num_of_scales, roi.width, CV_32F, this->gauss_corr_res.hostMem());
-        this->fw_all = cv::Mat(roi.height * num_of_feats, roi.width, CV_32F, this->data_features.hostMem());
-#else
-        uint width_freq = roi.width;
+    MatScaleFeats patch_feats{num_scales, num_features, roi};
+    MatScaleFeats temp{num_scales, num_features, roi};
 
-        this->in_all = cv::Mat(roi, CV_32F);
-#endif
+    KCF_Tracker::GaussianCorrelation gaussian_correlation{num_scales, num_features, roi};
 
-        this->data_i_features = DynMem(cells_size * num_of_feats);
-        this->data_i_1ch = DynMem(cells_size * num_of_scales);
+    MatScales ifft2_res{num_scales, roi};
 
-        this->ifft2_res = cv::Mat(roi, CV_32FC(num_of_feats), this->data_i_features.hostMem());
-        this->response = cv::Mat(roi, CV_32FC(num_of_scales), this->data_i_1ch.hostMem());
+    ComplexMat zf{uint(freq_size.height), uint(freq_size.width), num_features, num_scales};
+    ComplexMat kzf{uint(freq_size.height), uint(freq_size.width), 1, num_scales};
 
-#ifdef CUFFT
-        this->zf.create(roi.height, width_freq, num_of_feats, num_of_scales, this->stream);
-        this->kzf.create(roi.height, width_freq, num_of_scales, this->stream);
-        this->kf.create(roi.height, width_freq, num_of_scales, this->stream);
-#else
-        this->zf.create(roi.height, width_freq, num_of_feats, num_of_scales);
-        this->kzf.create(roi.height, width_freq, num_of_scales);
-        this->kf.create(roi.height, width_freq, num_of_scales);
-#endif
-
-#ifdef BIG_BATCH
-        if (num_of_scales > 1) {
-            this->max_responses.reserve(num_of_scales);
-            this->max_locs.reserve(num_of_scales);
-            this->response_maps.reserve(num_of_scales);
-        }
-#endif
-    }
-    ThreadCtx(ThreadCtx &&) = default;
-    ~ThreadCtx()
-    {
-#if  !defined(BIG_BATCH) && defined(CUFFT) && (defined(ASYNC) || defined(OPENMP))
-        CudaSafeCall(cudaStreamDestroy(this->stream));
-#endif
-    }
-
-    const double scale;
+public:
 #ifdef ASYNC
     std::future<void> async_res;
 #endif
 
-    DynMem xf_sqr_norm, yf_sqr_norm;
-
-    cv::Mat in_all, fw_all, ifft2_res, response;
-    ComplexMat zf, kzf, kf, xyf;
-
-    DynMem data_i_features, data_i_1ch;
-    // CuFFT and FFTW variables
-    DynMem gauss_corr_res, data_features;
+    MatScales response{num_scales, roi};
 
-    // CuFFT variables
-    cudaStream_t stream = nullptr;
-    ComplexMat model_alphaf, model_xf;
-
-    // Variables used during non big batch mode and in big batch mode with ThreadCtx in p_threadctxs in kcf  on zero index.
-    cv::Point2i max_loc;
-    double max_val, max_response;
+    struct Max {
+        cv::Point2i loc;
+        double response;
+    };
 
 #ifdef BIG_BATCH
-    // Stores value of responses, location of maximal response and response maps for each scale
-    std::vector<double> max_responses;
-    std::vector<cv::Point2i> max_locs;
-    std::vector<cv::Mat> response_maps;
+    std::vector<Max> max = std::vector<Max>(num_scales);
+#else
+    Max max;
+    const double scale;
 #endif
 };