1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70
|
#ifndef CAFFE2_OPERATORS_UTILS_NMS_GPU_H_
#define CAFFE2_OPERATORS_UTILS_NMS_GPU_H_
#include <vector>
#include "caffe2/core/context_gpu.h"
namespace caffe2 {
namespace utils {
// Computes Non-Maximum Suppression on the GPU
// Reject a bounding box if its region has an intersection-overunion (IoU)
// overlap with a higher scoring selected bounding box larger than a
// threshold.
//
// d_desc_sorted_boxes : pixel coordinates of proposed bounding boxes
// size: (N,4), format: [x1; y1; x2; y2]
// the boxes are sorted by scores in descending order
// N : number of boxes
// d_keep_sorted_list : row indices of the selected proposals, sorted by score
// h_nkeep : number of selected proposals
// dev_delete_mask, host_delete_mask : Tensors that will be used as temp storage
// by NMS
// Those tensors will be resized to the necessary size
// context : current CUDA context
TORCH_API void nms_gpu_upright(
const float* d_desc_sorted_boxes,
const int N,
const float thresh,
const bool legacy_plus_one,
int* d_keep_sorted_list,
int* h_nkeep,
TensorCUDA& dev_delete_mask,
TensorCPU& host_delete_mask,
CUDAContext* context);
struct RotatedBox {
float x_ctr, y_ctr, w, h, a;
};
// Same as nms_gpu_upright, but for rotated boxes with angle info.
// d_desc_sorted_boxes : pixel coordinates of proposed bounding boxes
// size: (N,5), format: [x_ct; y_ctr; width; height; angle]
// the boxes are sorted by scores in descending order
TORCH_API void nms_gpu_rotated(
const float* d_desc_sorted_boxes,
const int N,
const float thresh,
int* d_keep_sorted_list,
int* h_nkeep,
TensorCUDA& dev_delete_mask,
TensorCPU& host_delete_mask,
CUDAContext* context);
TORCH_API void nms_gpu(
const float* d_desc_sorted_boxes,
const int N,
const float thresh,
const bool legacy_plus_one,
int* d_keep_sorted_list,
int* h_nkeep,
TensorCUDA& dev_delete_mask,
TensorCPU& host_delete_mask,
CUDAContext* context,
const int box_dim);
} // namespace utils
} // namespace caffe2
#endif // CAFFE2_OPERATORS_UTILS_NMS_GPU_H_
|