src/complexmat.hpp

   1 #ifndef COMPLEX_MAT_HPP_213123048309482094
   2 #define COMPLEX_MAT_HPP_213123048309482094
   3
   4 #include <opencv2/opencv.hpp>
   5 #include <vector>
   6 #include <algorithm>
   7 #include <functional>
   8
   9 #ifdef CUFFT
  10   #include "managed_allocator.h"
  11
  12   template<class T>
  13   using managed_vector = std::vector<T,managed_allocator<T>>;
  14 #endif
  15
  16 template<typename T> class ComplexMat_
  17 {
  18 public:
  19     int cols;
  20     int rows;
  21     int n_channels;
  22     int n_scales = 1;
  23
  24     ComplexMat_() : cols(0), rows(0), n_channels(0) {}
  25     ComplexMat_(int _rows, int _cols, int _n_channels) : cols(_cols), rows(_rows), n_channels(_n_channels)
  26     {
  27         p_data.resize(n_channels*cols*rows);
  28     }
  29
  30
  31     //assuming that mat has 2 channels (real, img)
  32     ComplexMat_(const cv::Mat & mat) : cols(mat.cols), rows(mat.rows), n_channels(1)
  33     {
  34         p_data = convert(mat);
  35     }
  36
  37     ComplexMat_(int _rows, int _cols, int _n_channels,std::vector<std::complex<T>> data) : cols(_cols), rows(_rows), n_channels(_n_channels)
  38     {
  39         p_data = data;
  40     }
  41
  42     void create(int _rows, int _cols, int _n_channels)
  43     {
  44         rows = _rows;
  45         cols = _cols;
  46         n_channels = _n_channels;
  47         p_data.resize(n_channels*cols*rows);
  48     }
  49
  50     void create(int _rows, int _cols, int _n_channels, int _n_scales)
  51     {
  52         rows = _rows;
  53         cols = _cols;
  54         n_channels = _n_channels;
  55         n_scales = _n_scales;
  56         p_data.resize(n_channels*cols*rows);
  57     }
  58     // cv::Mat API compatibility
  59     cv::Size size() { return cv::Size(cols, rows); }
  60     int channels() { return n_channels; }
  61     int channels() const { return n_channels; }
  62
  63     //assuming that mat has 2 channels (real, imag)
  64     void set_channel(int idx, const cv::Mat & mat)
  65     {
  66         assert(idx >= 0 && idx < n_channels);
  67         for (int i = 0; i < rows; ++i){
  68             const std::complex<T> *row = mat.ptr<std::complex<T>>(i);
  69             for (int j = 0; j < cols; ++j)
  70                 p_data[idx*rows*cols+i*cols+j]=row[j];
  71         }
  72     }
  73
  74
  75     std::vector<T> sqr_norm() const
  76     {
  77         std::vector<T> sums_sqr_norms;
  78         sums_sqr_norms.resize(n_scales);
  79         int n_channels_per_scale = n_channels/n_scales;
  80         int scale_offset = n_channels_per_scale*rows*cols;
  81         T sum_sqr_norm;
  82         for(int scale = 0; scale < n_scales; ++scale){
  83             sum_sqr_norm = 0;
  84             for (int i = 0; i < n_channels_per_scale; ++i)
  85                 for (auto lhs = p_data.begin()+i*rows*cols+scale*scale_offset; lhs != p_data.begin()+(i+1)*rows*cols+scale*scale_offset; ++lhs)
  86                     sum_sqr_norm += lhs->real()*lhs->real() + lhs->imag()*lhs->imag();
  87             sums_sqr_norms[scale] = sum_sqr_norm/static_cast<T>(cols*rows);
  88         }
  89         return sums_sqr_norms;
  90     }
  91
  92     ComplexMat_<T> sqr_mag() const
  93     {
  94         return mat_const_operator( [](std::complex<T> & c) { c = c.real()*c.real() + c.imag()*c.imag(); } );
  95     }
  96
  97     ComplexMat_<T> conj() const
  98     {
  99         return mat_const_operator( [](std::complex<T> & c) { c = std::complex<T>(c.real(), -c.imag()); } );
 100     }
 101
 102     ComplexMat_<T> sum_over_channels() const
 103     {
 104         assert(p_data.size() > 1);
 105         ComplexMat_<T> result(this->rows, this->cols, 1);
 106         std::copy(p_data.begin(),p_data.begin()+rows*cols, result.p_data.begin());
 107         for (int i = 1; i < n_channels; ++i) {
 108             std::transform(result.p_data.begin(), result.p_data.end(), p_data.begin()+i*rows*cols, result.p_data.begin(), std::plus<std::complex<T>>());
 109         }
 110         return result;
 111     }
 112
 113     //return 2 channels (real, imag) for first complex channel
 114     cv::Mat to_cv_mat() const
 115     {
 116         assert(p_data.size() >= 1);
 117         return channel_to_cv_mat(0);
 118     }
 119     // return a vector of 2 channels (real, imag) per one complex channel
 120     std::vector<cv::Mat> to_cv_mat_vector() const
 121     {
 122         std::vector<cv::Mat> result;
 123         result.reserve(n_channels);
 124
 125         for (int i = 0; i < n_channels; ++i)
 126             result.push_back(channel_to_cv_mat(i));
 127
 128         return result;
 129     }
 130
 131     std::complex<T>* get_p_data() const
 132     {
 133         return p_data.data();
 134     }
 135
 136     ComplexMat_<T> get_part(int id, int n_of_feat)
 137     {
 138         std::vector<std::complex<T>> data(p_data.begin()+id*rows*cols*n_of_feat,
 139                                       p_data.begin()+(id+1)*rows*cols*n_of_feat);
 140         ComplexMat_<T> result(this->rows,this->cols,n_of_feat,data);
 141         return result;
 142     }
 143
 144     //element-wise per channel multiplication, division and addition
 145     ComplexMat_<T> operator*(const ComplexMat_<T> & rhs) const
 146     {
 147         return mat_mat_operator( [](std::complex<T> & c_lhs, const std::complex<T> & c_rhs) { c_lhs *= c_rhs; }, rhs);
 148     }
 149     ComplexMat_<T> operator/(const ComplexMat_<T> & rhs) const
 150     {
 151         return mat_mat_operator( [](std::complex<T> & c_lhs, const std::complex<T> & c_rhs) { c_lhs /= c_rhs; }, rhs);
 152     }
 153     ComplexMat_<T> operator+(const ComplexMat_<T> & rhs) const
 154     {
 155         return mat_mat_operator( [](std::complex<T> & c_lhs, const std::complex<T> & c_rhs)  { c_lhs += c_rhs; }, rhs);
 156     }
 157
 158     //multiplying or adding constant
 159     ComplexMat_<T> operator*(const T & rhs) const
 160     {
 161         return mat_const_operator( [&rhs](std::complex<T> & c) { c *= rhs; });
 162     }
 163     ComplexMat_<T> operator+(const T & rhs) const
 164     {
 165         return mat_const_operator( [&rhs](std::complex<T> & c) { c += rhs; });
 166     }
 167
 168     //multiplying element-wise multichannel by one channel mats (rhs mat is with one channel)
 169     ComplexMat_<T> mul(const ComplexMat_<T> & rhs) const
 170     {
 171         return matn_mat1_operator( [](std::complex<T> & c_lhs, const std::complex<T> & c_rhs) { c_lhs *= c_rhs; }, rhs);
 172     }
 173
 174     //multiplying element-wise multichannel by one channel mats (rhs mat is with multiple channel)
 175     ComplexMat_<T> mul2(const ComplexMat_<T> & rhs) const
 176     {
 177         return matn_mat2_operator( [](std::complex<T> & c_lhs, const std::complex<T> & c_rhs) { c_lhs *= c_rhs; }, rhs);
 178     }
 179
 180     //text output
 181     friend std::ostream & operator<<(std::ostream & os, const ComplexMat_<T> & mat)
 182     {
 183         //for (int i = 0; i < mat.n_channels; ++i){
 184         for (int i = 0; i < 1; ++i){
 185             os << "Channel " << i << std::endl;
 186             for (int j = 0; j < mat.rows; ++j) {
 187                 for (int k = 0; k < mat.cols-1; ++k)
 188                     os << mat.p_data[j*mat.cols + k] << ", ";
 189                 os << mat.p_data[j*mat.cols + mat.cols-1] << std::endl;
 190             }
 191         }
 192         return os;
 193     }
 194
 195
 196 private:
 197 #ifdef CUFFT
 198     mutable managed_vector<std::complex<T>> p_data;
 199 #else
 200     mutable std::vector<std::complex<T>> p_data;
 201 #endif
 202     //convert 2 channel mat (real, imag) to vector row-by-row
 203     std::vector<std::complex<T>> convert(const cv::Mat & mat)
 204     {
 205         std::vector<std::complex<T>> result;
 206         result.reserve(mat.cols*mat.rows);
 207         for (int y = 0; y < mat.rows; ++y) {
 208             const T * row_ptr = mat.ptr<T>(y);
 209             for (int x = 0; x < 2*mat.cols; x += 2){
 210                 result.push_back(std::complex<T>(row_ptr[x], row_ptr[x+1]));
 211             }
 212         }
 213         return result;
 214     }
 215
 216     ComplexMat_<T> mat_mat_operator(void (*op)(std::complex<T> & c_lhs, const std::complex<T> & c_rhs), const ComplexMat_<T> & mat_rhs) const
 217     {
 218         assert(mat_rhs.n_channels == n_channels && mat_rhs.cols == cols && mat_rhs.rows == rows);
 219
 220         ComplexMat_<T> result = *this;
 221         for (int i = 0; i < n_channels; ++i) {
 222             auto lhs = result.p_data.begin()+i*rows*cols;
 223             auto rhs = mat_rhs.p_data.begin()+i*rows*cols;
 224             for ( ; lhs != result.p_data.begin()+(i+1)*rows*cols; ++lhs, ++rhs)
 225                 op(*lhs, *rhs);
 226         }
 227
 228         return result;
 229     }
 230     ComplexMat_<T> matn_mat1_operator(void (*op)(std::complex<T> & c_lhs, const std::complex<T> & c_rhs), const ComplexMat_<T> & mat_rhs) const
 231     {
 232         assert(mat_rhs.n_channels == 1 && mat_rhs.cols == cols && mat_rhs.rows == rows);
 233
 234         ComplexMat_<T> result = *this;
 235         for (int i = 0; i < n_channels; ++i) {
 236             auto lhs = result.p_data.begin()+i*rows*cols;
 237             auto rhs = mat_rhs.p_data.begin();
 238             for ( ; lhs != result.p_data.begin()+(i+1)*rows*cols; ++lhs, ++rhs)
 239                 op(*lhs, *rhs);
 240         }
 241
 242         return result;
 243     }
 244     ComplexMat_<T> matn_mat2_operator(void (*op)(std::complex<T> & c_lhs, const std::complex<T> & c_rhs), const ComplexMat_<T> & mat_rhs) const
 245     {
 246         assert(mat_rhs.n_channels == n_channels/n_scales && mat_rhs.cols == cols && mat_rhs.rows == rows);
 247
 248         int n_channels_per_scale = n_channels/n_scales;
 249         int scale_offset = n_channels_per_scale*rows*cols;
 250         ComplexMat_<T> result = *this;
 251         for (int i = 0; i < n_scales; ++i) {
 252             for (int j = 0; j < n_channels_per_scale; ++j) {
 253                 auto lhs = result.p_data.begin()+(j*rows*cols)+(i*scale_offset);
 254                 auto rhs = mat_rhs.p_data.begin()+(j*rows*cols);
 255                 for ( ; lhs != result.p_data.begin()+((j+1)*rows*cols)+(i*scale_offset); ++lhs, ++rhs)
 256                     op(*lhs, *rhs);
 257             }
 258         }
 259
 260         return result;
 261     }
 262     ComplexMat_<T> mat_const_operator(const std::function<void(std::complex<T> & c_rhs)> & op) const
 263     {
 264         ComplexMat_<T> result = *this;
 265         for (int i = 0; i < n_channels; ++i)
 266             for (auto lhs = result.p_data.begin()+i*rows*cols; lhs != result.p_data.begin()+(i+1)*rows*cols; ++lhs)
 267                 op(*lhs);
 268         return result;
 269     }
 270
 271     cv::Mat channel_to_cv_mat(int channel_id) const
 272     {
 273         cv::Mat result(rows, cols, CV_32FC2);
 274         for (int y = 0; y < rows; ++y) {
 275             std::complex<T> * row_ptr = result.ptr<std::complex<T>>(y);
 276             for (int x = 0; x < cols; ++x){
 277                 row_ptr[x] = p_data[channel_id*rows*cols+y*cols+x];
 278             }
 279         }
 280         return result;
 281     }
 282
 283 };
 284
 285 typedef ComplexMat_<float> ComplexMat;
 286
 287
 288 #endif //COMPLEX_MAT_HPP_213123048309482094