void roi_sampling_forward_impl()

in src/roi_sampling/roi_sampling_cpu.cpp [8:45]


void roi_sampling_forward_impl(
    at::TensorAccessor<scalar_t, 4> x,
    at::TensorAccessor<coord_t, 2> bbx,
    at::TensorAccessor<int64_t, 1> idx,
    at::TensorAccessor<scalar_t, 4> y,
    at::TensorAccessor<uint8_t, 3> mask,
    bool valid_mask,
    Sampler sampler) {
  auto roi_height = static_cast<coord_t>(y.size(2)),
       roi_width  = static_cast<coord_t>(y.size(3));
  auto img_height = static_cast<coord_t>(x.size(2)),
       img_width  = static_cast<coord_t>(x.size(3));

  for (int64_t n = 0; n < idx.size(0); ++n) {
    auto img_idx = idx[n];
    auto i0 = bbx[n][0], j0 = bbx[n][1], i1 = bbx[n][2], j1 = bbx[n][3];

    for (int64_t c = 0; c < x.size(1); ++c) {
      // Create indexer for this plane and image
      auto accessor = x[img_idx][c];

      for (int64_t i_roi = 0; i_roi < y.size(2); ++i_roi) {
        auto y_img = roi_to_img(static_cast<coord_t>(i_roi) + coord_t(0.5), i0, i1, roi_height);

        for (int64_t j_roi = 0; j_roi < y.size(3); ++j_roi) {
          auto x_img = roi_to_img(static_cast<coord_t>(j_roi) + coord_t(0.5), j0, j1, roi_width);

          y[n][c][i_roi][j_roi] = sampler.forward(y_img, x_img, accessor);

          // Optionally write to mask
          if (valid_mask) {
            mask[n][i_roi][j_roi] = y_img >= 0 && y_img < img_height && x_img >= 0 && x_img < img_width;
          }
        }
      }
    }
  }
}