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
CAFFE2_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
CAFFE2_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);
CAFFE2_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_
|
