at::Tensor nms_cpu()

in src/nms/nms_cpu.cpp [60:95]

• 23 lines of code
• 4 McCabe index (conditional complexity)

at::Tensor nms_cpu(const at::Tensor& comp_mat, const at::Tensor& idx, int n_max) {
  int64_t num = comp_mat.size(0);

  auto _comp_mat = comp_mat.accessor<int64_t, 2>();
  auto _idx = idx.data<int64_t>();

  // Copy to C++ data structures
  std::list<int64_t> candidates;
  std::copy(_idx, _idx + num, std::back_inserter(candidates));

  std::vector<int64_t> selection;
  size_t n_max_ = n_max > 0 ? n_max : num;

  // Run actual nms
  while (!candidates.empty() && selection.size() < n_max_) {
    // Select first element
    auto i = candidates.front();
    selection.push_back(i);
    candidates.pop_front();

    // Remove conflicts
    candidates.remove_if([&_comp_mat,&i] (const int64_t &j) {
      auto ii = std::min(i, j), jj = std::max(i, j);

      auto block_idx = jj / THREADS_PER_BLOCK;
      auto bit_idx = jj % THREADS_PER_BLOCK;
      return _comp_mat[ii][block_idx] & (int64_t(1) << bit_idx);
    });
  }

  // Copy to output
  auto selection_tensor = at::zeros(selection.size(), comp_mat.options());
  std::copy(selection.begin(), selection.end(), selection_tensor.data<int64_t>());

  return selection_tensor;
}