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