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