inference-engine/src/extension/ext_spatial_transformer.cpp

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