Extend to support not only multiple scales but also multiple angles

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

diff --git a/src/threadctx.hpp b/src/threadctx.hpp
index 982905b96699b5928d39eb2ba158fc6920c2120f..0b23707172cbbc1033e5de02980c076041f8834f 100644 (file)
--- a/src/threadctx.hpp
+++ b/src/threadctx.hpp
@@ -5,25 +5,56 @@
 #include "dynmem.hpp"
 #include "kcf.h"
 #include "complexmat.hpp"
+#include <vector>
 
 class KCF_Tracker;
 
+template <typename T>
+class ScaleRotVector : public std::vector<T> {
+public:
+    ScaleRotVector(const std::vector<double> &scales, const std::vector<double> &angles)
+        : scales(scales)
+        , angles(angles)
+    {}
+
+    uint getIdx(uint scale_idx, uint angle_idx) const { return angles.size() * scale_idx + angle_idx; }
+    uint getScaleIdx(uint idx) const { return idx / angles.size(); }
+    uint getAngleIdx(uint idx) const { return idx % angles.size(); }
+    T& operator()(uint scale_idx, uint angle_idx) { return std::vector<T>::at(getIdx(scale_idx, angle_idx)); }
+    double scale(uint idx) const { return scales[getScaleIdx(idx)]; }
+    double angle(uint idx) const { return angles[getAngleIdx(idx)]; }
+private:
+    const std::vector<double> scales, angles;
+};
+
 struct ThreadCtx {
   public:
     ThreadCtx(cv::Size roi, uint num_features
 #ifdef BIG_BATCH
-              , uint num_scales
+              , const std::vector<double> &scales
+              , const std::vector<double> &angles
 #else
               , double scale
+              , double angle
 #endif
              )
         : roi(roi)
         , num_features(num_features)
-        , num_scales(IF_BIG_BATCH(num_scales, 1))
-#ifndef BIG_BATCH
+        , num_scales(IF_BIG_BATCH(scales.size(), 1))
+        , num_angles(IF_BIG_BATCH(angles.size(), 1))
+#ifdef BIG_BATCH
+        , max(scales, angles)
+        , dbg_patch(scales, angles)
+        {
+            max.resize(scales.size() * angles.size());
+            dbg_patch.resize(scales.size() * angles.size());
+        }
+#else
         , scale(scale)
+        , angle(angle)
+        {}
 #endif
-    {}
+
 
     ThreadCtx(ThreadCtx &&) = default;
 
@@ -32,24 +63,25 @@ private:
     cv::Size roi;
     uint num_features;
     uint num_scales;
+    uint num_angles;
     cv::Size freq_size = Fft::freq_size(roi);
 
-    MatScaleFeats patch_feats{num_scales, num_features, roi};
-    MatScaleFeats temp{num_scales, num_features, roi};
+    MatScaleFeats patch_feats{num_scales * num_angles, num_features, roi};
+    MatScaleFeats temp{num_scales * num_angles, num_features, roi};
 
-    KCF_Tracker::GaussianCorrelation gaussian_correlation{num_scales, num_features, roi};
+    KCF_Tracker::GaussianCorrelation gaussian_correlation{num_scales * num_angles, num_features, roi};
 
-    MatScales ifft2_res{num_scales, roi};
+    MatScales ifft2_res{num_scales * num_angles, roi};
 
-    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};
+    ComplexMat zf{uint(freq_size.height), uint(freq_size.width), num_features, num_scales * num_angles};
+    ComplexMat kzf{uint(freq_size.height), uint(freq_size.width), 1, num_scales * num_angles};
 
 public:
 #ifdef ASYNC
     std::future<void> async_res;
 #endif
 
-    MatScales response{num_scales, roi};
+    MatScales response{num_scales * num_angles, roi};
 
     struct Max {
         cv::Point2i loc;
@@ -57,10 +89,12 @@ public:
     };
 
 #ifdef BIG_BATCH
-    std::vector<Max> max = std::vector<Max>(num_scales);
+    ScaleRotVector<Max> max;
+    ScaleRotVector<cv::Mat> dbg_patch; // images for visual debugging
 #else
     Max max;
-    const double scale;
+    const double scale, angle;
+    cv::Mat dbg_patch; // image for visual debugging
 #endif
 };