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