operator/contrib/rroi_align.cc

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file rroi_align.cc
 * \brief rroi align operator
 * \author Yixin Bao
 * Forward pass adapted from Caffe2
 * link: https://github.com/pytorch/pytorch/blob/master/caffe2/operators/roi_align_rotated_op.cc
 */
#include "./rroi_align-inl.h"
#include <mshadow/tensor.h>
#include "math.h"

using std::max;
using std::min;
using std::floor;
using std::ceil;

namespace mxnet {
namespace op {

template <typename DType>
struct position_for_bilinear_interpolate {
  // 4 positions and corresponding weights for
  // computing bilinear interpolation
  int pos1, pos2, pos3, pos4;
  DType w1, w2, w3, w4;
};

template <typename DType>
void pre_calc_for_bilinear_interpolate(
    const int height, const int width, const int pooled_height, const int pooled_width,
    const int iy_upper, const int ix_upper, DType roi_start_h, DType roi_start_w,
    DType bin_size_h, DType bin_size_w, int roi_bin_grid_h, int roi_bin_grid_w,
    DType roi_center_h, DType roi_center_w, DType theta,
    std::vector<position_for_bilinear_interpolate<DType>> *pre_calc) {
  int pre_calc_index = 0;
  DType cosTheta = cos(theta);
  DType sinTheta = sin(theta);
  for (int ph = 0; ph < pooled_height; ph++) {
    for (int pw = 0; pw < pooled_width; pw++) {
      // calc bin grid position (xx,yy)
      for (int iy = 0; iy < iy_upper; iy++) {
        const DType yy = roi_start_h + ph * bin_size_h +
            static_cast<DType>(iy + .5f) * bin_size_h /
                static_cast<DType>(roi_bin_grid_h);  // e.g., 0.5, 1.5
        for (int ix = 0; ix < ix_upper; ix++) {
          const DType xx = roi_start_w + pw * bin_size_w +
              static_cast<DType>(ix + .5f) * bin_size_w /
                  static_cast<DType>(roi_bin_grid_w);

          // Rotate by theta around the center and translate
          DType x = xx * cosTheta + yy * sinTheta + roi_center_w;
          DType y = yy * cosTheta - xx * sinTheta + roi_center_h;

          // deal with: inverse elements are out of feature map boundary
          if (y < -1.0 || y > height || x < -1.0 || x > width) {
            // empty
            position_for_bilinear_interpolate<DType> &pc = (*pre_calc)[pre_calc_index];
            pc.pos1 = 0;
            pc.pos2 = 0;
            pc.pos3 = 0;
            pc.pos4 = 0;
            pc.w1 = 0;
            pc.w2 = 0;
            pc.w3 = 0;
            pc.w4 = 0;
            pre_calc_index += 1;
            continue;
          }
          if (y <= 0) {
            y = 0;
          }
          if (x <= 0) {
            x = 0;
          }

          // calc 4 points for interpolation
          int y_low = static_cast<int>(y);
          int x_low = static_cast<int>(x);
          int y_high;
          int x_high;
          if (y_low >= height - 1) {
            y_high = y_low = height - 1;
            y = (DType)y_low;
          } else {
            y_high = y_low + 1;
          }
          if (x_low >= width - 1) {
            x_high = x_low = width - 1;
            x = (DType)x_low;
          } else {
            x_high = x_low + 1;
          }
          DType ly = y - y_low;
          DType lx = x - x_low;
          DType hy = 1. - ly, hx = 1. - lx;
          DType w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;

          // Save weights and indices
          position_for_bilinear_interpolate<DType> &pc = (*pre_calc)[pre_calc_index];
          pc.pos1 = y_low * width + x_low;
          pc.pos2 = y_low * width + x_high;
          pc.pos3 = y_high * width + x_low;
          pc.pos4 = y_high * width + x_high;
          pc.w1 = w1;
          pc.w2 = w2;
          pc.w3 = w3;
          pc.w4 = w4;
          pre_calc_index += 1;
        }
      }
    }
  }
}

template <typename DType>
inline void RROIAlignForward(const OpContext &ctx, const RROIAlignParam &param,
                             const std::vector<TBlob> &in_data, const std::vector<OpReqType> &req,
                             const std::vector<TBlob> &out_data) {
  // data: [batch_size, c, h, w]
  const TBlob &data = in_data[rroialign::kData];
  const TBlob &bbox = in_data[rroialign::kBox];
  const DType *bottom_data = data.dptr<DType>();
  const int channels_ = data.size(1);
  const int height_ = data.size(2);
  const int width_ = data.size(3);
  const index_t data_size_c = height_ * width_;
  const index_t data_size = channels_ * data_size_c;

  // bbox: [num_rois, 6] (6: [batch_index, x, y, w, h, theta])
  const DType *bottom_rois = bbox.dptr<DType>();
  const int num_rois = bbox.size(0);
  const float spatial_scale_ = param.spatial_scale;
  const int sampling_ratio_ = param.sampling_ratio;

  // out: [num_rois, c, pooled_h, pooled_w]
  const TBlob &out = out_data[rroialign::kOut];
  DType *top_data = out.dptr<DType>();
  const int pooled_height_ = out.size(2);
  const int pooled_width_ = out.size(3);
  const index_t out_size_c = pooled_height_ * pooled_width_;
  const index_t out_size = channels_ * out_size_c;

  // (n, c, ph, pw) is an element in the pooled output
  // can be parallelized using omp
#pragma omp parallel for num_threads(engine::OpenMP::Get()->GetRecommendedOMPThreadCount())
  for (int n = 0; n < num_rois; ++n) {
    // Increment ROI data pointer
    const DType *bottom_rois_n = bottom_rois + n * bbox.size(1);
    DType *top_data_n = top_data + n * out_size;
    int roi_batch_ind = static_cast<int>(bottom_rois_n[0]);
    DType roi_center_w = bottom_rois_n[1] * spatial_scale_;
    DType roi_center_h = bottom_rois_n[2] * spatial_scale_;
    DType roi_width = bottom_rois_n[3] * spatial_scale_;
    DType roi_height = bottom_rois_n[4] * spatial_scale_;
    DType roi_theta = bottom_rois_n[5] * M_PI / 180.0;

    // force malformed ROIs to be 1 * 1
    roi_width = max(roi_width, (DType) 1.);
    roi_height = max(roi_height, (DType) 1.);
    // roi_start_h and roi_start_w are computed wrt the center of RoI (x, y).
    // Appropriate translation needs to be applied after.
    DType roi_start_h = -roi_height / 2.0;
    DType roi_start_w = -roi_width / 2.0;

    const DType bin_size_h = static_cast<DType>(roi_height) / static_cast<DType>(pooled_height_);
    const DType bin_size_w = static_cast<DType>(roi_width) / static_cast<DType>(pooled_width_);
    // We use roi_bin_grid to sample the grid and mimic integral,
    // e.g. roi_bin_grid = 2, means sample 2*2=4 points in each bin
    int roi_bin_grid_h =
        (sampling_ratio_ > 0) ? sampling_ratio_ : ceil(roi_height / pooled_height_);
    int roi_bin_grid_w = (sampling_ratio_ > 0) ? sampling_ratio_ : ceil(roi_width / pooled_width_);
    const DType bin_points_count = roi_bin_grid_h * roi_bin_grid_w;  // e.g. = 4

    // We want to precalculate indices and weights shared by all channels,
    // this is the key point of optimization.
    std::vector<position_for_bilinear_interpolate<DType>> pre_calc(roi_bin_grid_h * roi_bin_grid_w *
                                                                   pooled_width_ * pooled_height_);

    pre_calc_for_bilinear_interpolate(height_, width_, pooled_height_, pooled_width_,
                                      roi_bin_grid_h, roi_bin_grid_w, roi_start_h, roi_start_w,
                                      bin_size_h, bin_size_w, roi_bin_grid_h, roi_bin_grid_w,
                                      roi_center_h, roi_center_w, roi_theta, &pre_calc);

    for (int c = 0; c < channels_; ++c) {
      const DType *offset_bottom_data = bottom_data + roi_batch_ind * data_size + c * data_size_c;
      int pre_calc_index = 0;

      for (int ph = 0; ph < pooled_height_; ph++) {
        for (int pw = 0; pw < pooled_width_; pw++) {
          DType output_val = 0.;
          for (int iy = 0; iy < roi_bin_grid_h; iy++) {
            for (int ix = 0; ix < roi_bin_grid_w; ix++) {
              position_for_bilinear_interpolate<DType> pc = pre_calc[pre_calc_index];
              output_val +=
                  pc.w1 * offset_bottom_data[pc.pos1] + pc.w2 * offset_bottom_data[pc.pos2] +
                  pc.w3 * offset_bottom_data[pc.pos3] + pc.w4 * offset_bottom_data[pc.pos4];

              pre_calc_index += 1;
            }
          }
          output_val /= bin_points_count;  // avg pooling for bin grid
          int index = c * pooled_height_ * pooled_width_ + ph * pooled_width_ + pw;
          top_data_n[index] = output_val;
        }   // for pw
      }   // for ph
    }   // for c
  }   // for n
}

template<typename xpu>
void RROIAlignForwardCompute(const nnvm::NodeAttrs& attrs,
                      const OpContext& ctx, const std::vector<TBlob>& in_data,
                      const std::vector<OpReqType>& req,
                      const std::vector<TBlob>& out_data) {
  const RROIAlignParam& param = nnvm::get<RROIAlignParam>(attrs.parsed);
  CHECK_EQ(in_data.size(), 2);
  CHECK_EQ(out_data.size(), 1);
  CHECK_EQ(out_data[rroialign::kOut].shape_[0], in_data[rroialign::kBox].shape_[0]);

  MSHADOW_REAL_TYPE_SWITCH(in_data[0].type_flag_, DType, {
    RROIAlignForward<DType>(ctx, param, in_data, req, out_data);
  })
}

template<typename xpu>
void RROIAlignBackwardCompute(const nnvm::NodeAttrs& attrs,
                      const OpContext& ctx, const std::vector<TBlob>& in_data,
                      const std::vector<OpReqType>& req,
                      const std::vector<TBlob>& out_data) {
  LOG(FATAL) << "RROIAlign: Backward is not supported.";
}

DMLC_REGISTER_PARAMETER(RROIAlignParam);

NNVM_REGISTER_OP(_contrib_RROIAlign)
.describe(R"code(Performs Rotated ROI Align on the input array.

This operator takes a 4D feature map as an input array and region proposals as `rois`,
then align the feature map over sub-regions of input and produces a fixed-sized output array.

Different from ROI Align, RROI Align uses rotated rois, which is suitable for text detection.
RRoIAlign computes the value of each sampling point by bilinear interpolation from the nearby
grid points on the rotated feature map. No quantization is performed on any coordinates
involved in the RoI, its bins, or the sampling points. Bilinear interpolation is used to
compute the exact values of the input features at four regularly sampled locations in
each RoI bin. Then the feature map can be aggregated by avgpooling.

References
----------

Ma, Jianqi, et al. "Arbitrary-Oriented Scene Text Detection via Rotation Proposals."
IEEE Transactions on Multimedia, 2018.

)code" ADD_FILELINE)
.set_num_inputs(2)
.set_num_outputs(1)
.set_attr<nnvm::FListInputNames>("FListInputNames",
    [](const NodeAttrs& attrs) {
  return std::vector<std::string>{"data", "rois"};
})
.set_attr<nnvm::FListOutputNames>("FListOutputNames",
    [](const NodeAttrs& attrs) {
  return std::vector<std::string>{"output"};
})
.set_attr_parser(ParamParser<RROIAlignParam>)
.set_attr<mxnet::FInferShape>("FInferShape", [](const nnvm::NodeAttrs& attrs,
      mxnet::ShapeVector *in_shape, mxnet::ShapeVector *out_shape){
  using namespace mshadow;
  const RROIAlignParam& param = nnvm::get<RROIAlignParam>(attrs.parsed);
  CHECK_EQ(in_shape->size(), 2U) << "Input:[data, rois]";
  // data: [batch_size, c, h, w]
  mxnet::TShape dshape = in_shape->at(rroialign::kData);
  CHECK_EQ(dshape.ndim(), 4U) << "data should be a 4D tensor";
  // bbox: [num_rois, 6]
  mxnet::TShape bshape = in_shape->at(rroialign::kBox);
  CHECK_EQ(bshape.ndim(), 2U) << "bbox should be a 2D tensor of shape [batch, 6]";
  CHECK_EQ(bshape[1], 6U) << "bbox should be a 2D tensor of shape [batch, 6]";
  // out: [num_rois, c, pooled_h, pooled_w]
  out_shape->clear();
  out_shape->push_back(Shape4(bshape[0], dshape[1], param.pooled_size[0], param.pooled_size[1]));
  return true;
})
.set_attr<nnvm::FInferType>("FInferType", [](const nnvm::NodeAttrs& attrs,
      std::vector<int> *in_type, std::vector<int> *out_type) {
  CHECK_EQ(in_type->size(), 2U);
  int dtype = (*in_type)[0];
  CHECK_EQ(dtype, (*in_type)[1]);
  CHECK_NE(dtype, -1) << "Input must have specified type";

  out_type->clear();
  out_type->push_back(dtype);
  return true;
})
.set_attr<FCompute>("FCompute<cpu>", RROIAlignForwardCompute<cpu>)
.add_argument("data", "NDArray-or-Symbol", "Input data to the pooling operator, a 4D Feature maps")
.add_argument("rois", "NDArray-or-Symbol", "Bounding box coordinates, a 2D array")
.add_arguments(RROIAlignParam::__FIELDS__());

NNVM_REGISTER_OP(_backward_RROIAlign)
.set_num_outputs(2)
.set_attr<nnvm::TIsBackward>("TIsBackward", true)
.set_attr_parser(ParamParser<RROIAlignParam>)
.set_attr<FCompute>("FCompute<cpu>", RROIAlignBackwardCompute<cpu>);

}  // namespace op
}  // namespace mxnet