Nms.h

Go to the documentation of this file.

// ----------------------------------------------------------------------------
// -                        Open3D: www.open3d.org                            -
// ----------------------------------------------------------------------------
// Copyright (c) 2018-2026 www.open3d.org
// SPDX-License-Identifier: MIT
// ----------------------------------------------------------------------------
//
// Reference:
// https://github.com/open-mmlab/OpenPCDet/blob/master/pcdet/ops/iou3d_nms/src/iou3d_nms_kernel.cu
//
// Reference:
// https://github.com/open-mmlab/mmdetection3d/blob/master/mmdet3d/ops/iou3d/src/iou3d_kernel.cu
// 3D IoU Calculation and Rotated NMS(modified from 2D NMS written by others)
// Written by Shaoshuai Shi
// All Rights Reserved 2019-2020.
 
#pragma once
 
#include <algorithm>
#include <cstdint>
#include <numeric>
#include <vector>
 
#include "open3d/core/CUDAUtils.h"
 
#ifdef BUILD_SYCL_MODULE
// NmsSYCLKernel takes a real sycl::queue&, so any TU that sees this
// declaration needs the full SYCL runtime type (not -fsycl compilation --
// NmsOps.cpp already gets this transitively via <c10/xpu/XPUDeviceProp.h>,
// but Nms.h is included before that, so it must include the header directly
// to avoid an incomplete-type error). NmsSYCL.cpp (built with -fsycl) is
// where the actual device kernels are compiled.
#include <sycl/sycl.hpp>
#endif
 
#include "open3d/ml/contrib/IoUImpl.h"
#include "open3d/utility/Helper.h"
 
namespace open3d {
namespace ml {
namespace contrib {
 
#ifdef BUILD_CUDA_MODULE
 
int NmsCUDAKernel(const float *boxes,
                  const float *scores,
                  int n,
                  double nms_overlap_thresh,
                  int64_t *keep_indices_out);
#endif
 
#ifdef BUILD_SYCL_MODULE
 
int NmsSYCLKernel(sycl::queue &queue,
                  const float *boxes,
                  const float *scores,
                  int n,
                  double nms_overlap_thresh,
                  int64_t *keep_indices_out);
#endif
 
std::vector<int64_t> NmsCPUKernel(const float *boxes,
                                  const float *scores,
                                  int n,
                                  double nms_overlap_thresh);
 
template <typename T>
inline std::vector<int64_t> SortIndexesDescending(const T *values,
                                                  int64_t num) {
    std::vector<int64_t> indices(num);
    std::iota(indices.begin(), indices.end(), 0);
    std::stable_sort(
            indices.begin(), indices.end(),
            [&values](int64_t i, int64_t j) { return values[i] > values[j]; });
    return indices;
}
 
OPEN3D_HOST_DEVICE inline int NmsGreedyKeepCore(const uint64_t *mask,
                                                const int64_t *sort_indices,
                                                int n,
                                                uint64_t *remv,
                                                int64_t *keep_indices) {
    // Avoid utility::DivUp() (host-only, uses std::div) so this function
    // stays callable from CUDA/SYCL device code.
    const int num_block_cols = (n + NMS_BLOCK_SIZE - 1) / NMS_BLOCK_SIZE;
    for (int j = 0; j < num_block_cols; j++) {
        remv[j] = 0;
    }
    int count = 0;
    for (int i = 0; i < n; i++) {
        int block_col_idx = i / NMS_BLOCK_SIZE;
        int inner_block_col_idx = i % NMS_BLOCK_SIZE;
 
        if (!(remv[block_col_idx] & (1ULL << inner_block_col_idx))) {
            keep_indices[count++] = sort_indices[i];
 
            const uint64_t *p = mask + i * num_block_cols;
            for (int j = block_col_idx; j < num_block_cols; j++) {
                remv[j] |= p[j];
            }
        }
    }
    return count;
}
 
inline std::vector<int64_t> NmsGreedyKeep(const uint64_t *mask,
                                          const int64_t *sort_indices,
                                          int n) {
    const int num_block_cols = utility::DivUp(n, NMS_BLOCK_SIZE);
    std::vector<uint64_t> remv(num_block_cols);
    std::vector<int64_t> keep_indices(n);
    int count = NmsGreedyKeepCore(mask, sort_indices, n, remv.data(),
                                  keep_indices.data());
    keep_indices.resize(count);
    return keep_indices;
}
 
}  // namespace contrib
}  // namespace ml
}  // namespace open3d