bbox_transform_op.cc

#include "bbox_transform_op.h"
#include "caffe2/operators/generate_proposals_op_util_boxes.h"

#ifdef CAFFE2_USE_MKL
#include "caffe2/mkl/operators/operator_fallback_mkl.h"
#endif // CAFFE2_USE_MKL

namespace caffe2 {
namespace {

REGISTER_CPU_OPERATOR(BBoxTransform, BBoxTransformOp<float, CPUContext>);

#ifdef CAFFE2_HAS_MKL_DNN
REGISTER_MKL_OPERATOR(
    BBoxTransform,
    mkl::MKLFallbackOp<BBoxTransformOp<float, CPUContext>>);
#endif // CAFFE2_HAS_MKL_DNN

// Input: box, delta Output: box
OPERATOR_SCHEMA(BBoxTransform)
    .NumInputs(3)
    .NumOutputs(1, 2)
    .SetDoc(R"DOC(
Transform proposal bounding boxes to target bounding box using bounding box
    regression deltas.
)DOC")
    .Arg("weights", "vector<float> weights [wx, wy, ww, wh] for the deltas")
    .Arg(
        "apply_scale",
        "bool (default true), transform the boxes to the scaled image space"
        " after applying the bbox deltas."
        "Set to false to match the detectron code, set to true for keypoint"
        " models and for backward compatibility")
    .Arg(
        "correct_transform_coords",
        "bool (default false), Correct bounding box transform coordates,"
        " see bbox_transform() in boxes.py "
        "Set to true to match the detectron code, set to false for backward"
        " compatibility")
    .Input(
        0,
        "rois",
        "Bounding box proposals in pixel coordinates, "
        "Size (M, 4), format [x1, y1, x2, y2], or"
        "Size (M, 5), format [batch_index, x1, y1, x2, y2]. "
        "If proposals from multiple images in a batch are present, they "
        "should be grouped sequentially and in incremental order.")
    .Input(
        1,
        "deltas",
        "bounding box translations and scales,"
        "size (M, 4*K), format [dx, dy, dw, dh], K = # classes")
    .Input(
        2,
        "im_info",
        "Image dimensions, size (batch_size, 3), "
        "format [img_height, img_width, img_scale]")
    .Output(
        0,
        "box_out",
        "Pixel coordinates of the transformed bounding boxes,"
        "Size (M, 4*K), format [x1, y1, x2, y2]")
    .Output(
        1,
        "roi_batch_splits",
        "Tensor of shape (batch_size) with each element denoting the number "
        "of RoIs belonging to the corresponding image in batch");

SHOULD_NOT_DO_GRADIENT(BBoxTransform);
} // namespace

template <>
bool BBoxTransformOp<float, CPUContext>::RunOnDevice() {
  const auto& roi_in = Input(0);
  const auto& delta_in = Input(1);
  const auto& iminfo_in = Input(2);
  auto* box_out = Output(0);

  const int N = roi_in.dim32(0);
  CAFFE_ENFORCE_EQ(roi_in.ndim(), 2);
  CAFFE_ENFORCE(roi_in.dim32(1) == 4 || roi_in.dim32(1) == 5);

  CAFFE_ENFORCE_EQ(delta_in.ndim(), 2);
  CAFFE_ENFORCE_EQ(delta_in.dim32(0), N);
  CAFFE_ENFORCE_EQ(delta_in.dim32(1) % 4, 0);
  const int num_classes = delta_in.dim32(1) / 4;

  CAFFE_ENFORCE_EQ(iminfo_in.ndim(), 2);
  CAFFE_ENFORCE_EQ(iminfo_in.dim32(1), 3);
  const int batch_size = iminfo_in.dim32(0);

  DCHECK_EQ(weights_.size(), 4);

  Eigen::Map<const ERArrXXf> boxes0(
      roi_in.data<float>(), roi_in.dim32(0), roi_in.dim32(1));
  Eigen::Map<const ERArrXXf> deltas0(
      delta_in.data<float>(), delta_in.dim32(0), delta_in.dim32(1));

  // Count the number of RoIs per batch
  vector<int> num_rois_per_batch(batch_size, 0);
  if (roi_in.dim32(1) == 4) {
    CAFFE_ENFORCE_EQ(batch_size, 1);
    num_rois_per_batch[0] = N;
  } else {
    const auto& roi_batch_ids = boxes0.col(0);
    for (int i = 0; i < roi_batch_ids.size(); ++i) {
      const int roi_batch_id = roi_batch_ids(i);
      CAFFE_ENFORCE_LT(roi_batch_id, batch_size);
      num_rois_per_batch[roi_batch_id]++;
    }
  }

  CAFFE_ENFORCE_EQ(iminfo_in.dims(), (vector<TIndex>{batch_size, 3}));
  Eigen::Map<const ERArrXXf> iminfo(
      iminfo_in.data<float>(), iminfo_in.dim(0), iminfo_in.dim(1));

  box_out->ResizeLike(delta_in);
  Eigen::Map<ERArrXXf> new_boxes(
      box_out->mutable_data<float>(), box_out->dim32(0), box_out->dim32(1));

  // We assume roi_in and delta_in over multiple batches are grouped
  // together in increasing order as generated by GenerateProposalsOp
  int offset = 0;
  for (int i = 0; i < batch_size; ++i) {
    const int num_rois = num_rois_per_batch[i];
    const auto& cur_iminfo = iminfo.row(i);
    const float scale_before = cur_iminfo(2);
    const float scale_after = apply_scale_ ? cur_iminfo(2) : 1.0;
    int img_h = int(cur_iminfo(0) / scale_before + 0.5);
    int img_w = int(cur_iminfo(1) / scale_before + 0.5);

    const auto& cur_boxes =
        boxes0.rightCols(4).block(offset, 0, num_rois, 4) / scale_before;
    for (int k = 0; k < num_classes; k++) {
      const auto& cur_deltas = deltas0.block(offset, k * 4, num_rois, 4);
      const auto& trans_boxes = utils::bbox_transform(
          cur_boxes,
          cur_deltas,
          weights_,
          utils::BBOX_XFORM_CLIP_DEFAULT,
          correct_transform_coords_);
      const auto& clip_boxes = utils::clip_boxes(trans_boxes, img_h, img_w);
      new_boxes.block(offset, k * 4, num_rois, 4) = clip_boxes * scale_after;
    }

    offset += num_rois;
  }

  if (OutputSize() > 1) {
    auto* roi_batch_splits = Output(1);
    roi_batch_splits->Resize(batch_size);
    Eigen::Map<EArrXf> roi_batch_splits_map(
        roi_batch_splits->mutable_data<float>(), batch_size);
    roi_batch_splits_map =
        Eigen::Map<const EArrXi>(num_rois_per_batch.data(), batch_size)
            .cast<float>();
  }

  return true;
}

} // namespace caffe2