inference-engine/src/extension/ext_spatial_transformer.cpp

   1 // Copyright (C) 2018 Intel Corporation
   2 //
   3 // SPDX-License-Identifier: Apache-2.0
   4 //
   5
   6 #include "ext_list.hpp"
   7 #include "ext_base.hpp"
   8
   9 #include "matrixmult.h"
  10
  11 #include <algorithm>
  12 #include <vector>
  13 #include <cmath>
  14 #include <map>
  15 #include <string>
  16
  17 namespace InferenceEngine {
  18 namespace Extensions {
  19 namespace Cpu {
  20
  21 class SpatialTransformerImpl: public ExtLayerBase {
  22 public:
  23     explicit SpatialTransformerImpl(const CNNLayer* layer) {
  24         try {
  25             if (layer->insData.size() != 2 || layer->outData.empty())
  26                 THROW_IE_EXCEPTION << "Incorrect number of input/output edges!";
  27
  28             addConfig(layer, {DataConfigurator(ConfLayout::PLN), DataConfigurator(ConfLayout::PLN)}, {DataConfigurator(ConfLayout::PLN)});
  29         } catch (InferenceEngine::details::InferenceEngineException &ex) {
  30             errorMsg = ex.what();
  31         }
  32     }
  33
  34     StatusCode execute(std::vector<Blob::Ptr>& inputs, std::vector<Blob::Ptr>& outputs,
  35                        ResponseDesc *resp) noexcept override {
  36         std::vector<size_t> real_dims = inputs[0]->getTensorDesc().getDims();
  37         size_t data_size = inputs[0]->size();
  38
  39         const auto *src_data = inputs[0]->cbuffer().as<const float *>();
  40         auto *theta = inputs[1]->buffer().as<float *>();
  41         auto *dst_data = outputs[0]->buffer().as<float *>();
  42
  43         auto N = real_dims[0];
  44         auto C = real_dims[1];
  45         auto output_H_ = real_dims[2];
  46         auto output_W_ = real_dims[3];
  47
  48         // Prepare input and output grid
  49         std::vector<float> input_grid_data(N * output_H_ * output_W_ * 2);
  50         std::vector<float> output_grid_data(3 * output_H_ * output_W_);
  51         for (int i = 0; i < output_H_ * output_W_; ++i) {
  52             output_grid_data[3 * i] = (i / output_W_) * 1.0 / output_H_ * 2 - 1;
  53             output_grid_data[3 * i + 1] = (i % output_W_) * 1.0 / output_W_ * 2 - 1;
  54             output_grid_data[3 * i + 2] = 1;
  55         }
  56
  57         // Actually execute
  58         for (int i = 0; i < N; ++i) {
  59             auto coordinates = input_grid_data.begin() + (output_H_ * output_W_ * 2) * i;
  60
  61             auto M_size = output_H_ * output_W_;
  62             auto N_size = 2;
  63             auto K_size = 3;
  64
  65             matrixMult(&output_grid_data[0], theta + 6 * i, &(*coordinates), M_size, N_size, K_size, true);
  66
  67             int row_idx;
  68             float px, py;
  69
  70             for (int j = 0; j < C; ++j) {
  71                 for (int s = 0; s < output_H_; ++s) {
  72                     for (int t = 0; t < output_W_; ++t) {
  73                         row_idx = output_W_ * s + t;
  74
  75                         px = coordinates[row_idx * 2];
  76                         py = coordinates[row_idx * 2 + 1];
  77
  78                         size_t dst_offset = ((i * C + j) * output_H_ + s) * output_W_ + t;
  79                         size_t src_offset = ((i * C + j) * output_H_ + 0) * output_W_ + 0;
  80                         dst_data[dst_offset] = transform_forward_cpu(src_data + src_offset, px, py);
  81                     }
  82                 }
  83             }
  84         }
  85         return OK;
  86     }
  87
  88 private:
  89     float transform_forward_cpu(const float *pic, float px, float py) {
  90         int H = 24;
  91         int W = 94;
  92
  93         float res = 0.0f;
  94         float x = (px + 1) / 2 * H;
  95         float y = (py + 1) / 2 * W;
  96
  97         int m, n;
  98         float w;
  99
 100         m = std::floor(x);
 101         n = std::floor(y);
 102         w = 0;
 103         if (m >= 0 && m < H && n >= 0 && n < W) {
 104             w = std::max<float>(0.0f, 1 - std::abs(x - m)) * std::max<float>(0.0f, 1 - std::abs(y - n));
 105             res += w * pic[m * W + n];
 106         }
 107
 108         m = std::floor(x) + 1;
 109         n = std::floor(y);
 110         w = 0;
 111         if (m >= 0 && m < H && n >= 0 && n < W) {
 112             w = std::max<float>(0.0f, 1 - std::abs(x - m)) * std::max<float>(0.0f, 1 - std::abs(y - n));
 113             res += w * pic[m * W + n];
 114         }
 115
 116         m = std::floor(x);
 117         n = std::floor(y) + 1;
 118         w = 0;
 119         if (m >= 0 && m < H && n >= 0 && n < W) {
 120             w = std::max<float>(0.0f, 1 - std::abs(x - m)) * std::max<float>(0.0f, 1 - std::abs(y - n));
 121             res += w * pic[m * W + n];
 122         }
 123
 124         m = std::floor(x) + 1;
 125         n = std::floor(y) + 1;
 126         w = 0;
 127         if (m >= 0 && m < H && n >= 0 && n < W) {
 128             w = std::max<float>(0.0f, 1 - std::abs(x - m)) * std::max<float>(0.0f, 1 - std::abs(y - n));
 129             res += w * pic[m * W + n];
 130         }
 131
 132         return res;
 133     }
 134 };
 135
 136 class SpatialTransformerShapeInfer : public IShapeInferImpl {
 137 public:
 138     StatusCode inferShapes(const std::vector<SizeVector>& inShapes,
 139                            const std::map<std::string, std::string>& params,
 140                            const std::map<std::string, Blob::Ptr>& blobs,
 141                            std::vector<SizeVector>& outShapes,
 142                            ResponseDesc* resp) noexcept override {
 143         outShapes.push_back(inShapes[0]);
 144         return InferenceEngine::OK;
 145     }
 146 };
 147
 148 REG_FACTORY_FOR(ImplFactory<SpatialTransformerImpl>, SpatialTransformer);
 149 REG_SHAPE_INFER_FOR_TYPE(SpatialTransformerShapeInfer, SpatialTransformer);
 150
 151 }  // namespace Cpu
 152 }  // namespace Extensions
 153 }  // namespace InferenceEngine