Caffe2 - C++ API
A deep learning, cross platform ML framework
generate_proposals_op_util_nms.h
1 // Copyright 2004-present Facebook. All Rights Reserved.
2 
3 #ifndef CAFFE2_OPERATORS_UTILS_NMS_H_
4 #define CAFFE2_OPERATORS_UTILS_NMS_H_
5 
6 #include <list>
7 #include <vector>
8 
9 #include "caffe2/utils/eigen_utils.h"
10 
11 #include "caffe2/core/logging.h"
12 #include "caffe2/utils/math.h"
13 
14 namespace caffe2 {
15 namespace utils {
16 
17 // Greedy non-maximum suppression for proposed bounding boxes
18 // Reject a bounding box if its region has an intersection-overunion (IoU)
19 // overlap with a higher scoring selected bounding box larger than a
20 // threshold.
21 // Reference: detectron/lib/utils/cython_nms.pyx
22 // proposals: pixel coordinates of proposed bounding boxes,
23 // size: (M, 4), format: [x1; y1; x2; y2]
24 // scores: scores for each bounding box, size: (M, 1)
25 // sorted_indices: indices that sorts the scores from high to low
26 // return: row indices of the selected proposals
27 template <class Derived1, class Derived2>
28 std::vector<int> nms_cpu(
29  const Eigen::ArrayBase<Derived1>& proposals,
30  const Eigen::ArrayBase<Derived2>& scores,
31  const std::vector<int>& sorted_indices,
32  float thresh,
33  int topN = -1) {
34  CAFFE_ENFORCE_EQ(proposals.rows(), scores.rows());
35  CAFFE_ENFORCE_EQ(proposals.cols(), 4);
36  CAFFE_ENFORCE_EQ(scores.cols(), 1);
37  CAFFE_ENFORCE_LE(sorted_indices.size(), proposals.rows());
38 
39  using EArrX = EArrXt<typename Derived1::Scalar>;
40 
41  auto x1 = proposals.col(0);
42  auto y1 = proposals.col(1);
43  auto x2 = proposals.col(2);
44  auto y2 = proposals.col(3);
45 
46  EArrX areas = (x2 - x1 + 1.0) * (y2 - y1 + 1.0);
47 
48  EArrXi order = AsEArrXt(sorted_indices);
49  std::vector<int> keep;
50  int ci = 0;
51  while (order.size() > 0) {
52  // exit if already enough proposals
53  if (topN >= 0 && keep.size() >= topN) {
54  break;
55  }
56 
57  int i = order[0];
58  keep.push_back(i);
59  ConstEigenVectorArrayMap<int> rest_indices(
60  order.data() + 1, order.size() - 1);
61  EArrX xx1 = GetSubArray(x1, rest_indices).cwiseMax(x1[i]);
62  EArrX yy1 = GetSubArray(y1, rest_indices).cwiseMax(y1[i]);
63  EArrX xx2 = GetSubArray(x2, rest_indices).cwiseMin(x2[i]);
64  EArrX yy2 = GetSubArray(y2, rest_indices).cwiseMin(y2[i]);
65 
66  EArrX w = (xx2 - xx1 + 1.0).cwiseMax(0.0);
67  EArrX h = (yy2 - yy1 + 1.0).cwiseMax(0.0);
68  EArrX inter = w * h;
69  EArrX ovr = inter / (areas[i] + GetSubArray(areas, rest_indices) - inter);
70 
71  // indices for sub array order[1:n]
72  auto inds = GetArrayIndices(ovr <= thresh);
73  order = GetSubArray(order, AsEArrXt(inds) + 1);
74  }
75 
76  return keep;
77 }
78 
79 // Greedy non-maximum suppression for proposed bounding boxes
80 // Reject a bounding box if its region has an intersection-overunion (IoU)
81 // overlap with a higher scoring selected bounding box larger than a
82 // threshold.
83 // Reference: detectron/lib/utils/cython_nms.pyx
84 // proposals: pixel coordinates of proposed bounding boxes,
85 // size: (M, 4), format: [x1; y1; x2; y2]
86 // scores: scores for each bounding box, size: (M, 1)
87 // return: row indices of the selected proposals
88 template <class Derived1, class Derived2>
89 std::vector<int> nms_cpu(
90  const Eigen::ArrayBase<Derived1>& proposals,
91  const Eigen::ArrayBase<Derived2>& scores,
92  float thres) {
93  std::vector<int> indices(proposals.rows());
94  std::iota(indices.begin(), indices.end(), 0);
95  std::sort(
96  indices.data(),
97  indices.data() + indices.size(),
98  [&scores](int lhs, int rhs) { return scores(lhs) > scores(rhs); });
99 
100  return nms_cpu(proposals, scores, indices, thres);
101 }
102 
118 template <class Derived1, class Derived2, class Derived3>
119 std::vector<int> soft_nms_cpu(
120  Eigen::ArrayBase<Derived3>* out_scores,
121  const Eigen::ArrayBase<Derived1>& proposals,
122  const Eigen::ArrayBase<Derived2>& scores,
123  const std::vector<int>& indices,
124  float sigma = 0.5,
125  float overlap_thresh = 0.3,
126  float score_thresh = 0.001,
127  unsigned int method = 1,
128  int topN = -1) {
129  CAFFE_ENFORCE_EQ(proposals.rows(), scores.rows());
130  CAFFE_ENFORCE_EQ(proposals.cols(), 4);
131  CAFFE_ENFORCE_EQ(scores.cols(), 1);
132 
133  using EArrX = EArrXt<typename Derived1::Scalar>;
134 
135  const auto& x1 = proposals.col(0);
136  const auto& y1 = proposals.col(1);
137  const auto& x2 = proposals.col(2);
138  const auto& y2 = proposals.col(3);
139 
140  EArrX areas = (x2 - x1 + 1.0) * (y2 - y1 + 1.0);
141 
142  // Initialize out_scores with original scores. Will be iteratively updated
143  // as Soft-NMS is applied.
144  *out_scores = scores;
145 
146  std::vector<int> keep;
147  EArrXi pending = AsEArrXt(indices);
148  while (pending.size() > 0) {
149  // Exit if already enough proposals
150  if (topN >= 0 && keep.size() >= topN) {
151  break;
152  }
153 
154  // Find proposal with max score among remaining proposals
155  int max_pos;
156  auto max_score = GetSubArray(*out_scores, pending).maxCoeff(&max_pos);
157  int i = pending[max_pos];
158  keep.push_back(i);
159 
160  // Compute IoU of the remaining boxes with the identified max box
161  std::swap(pending(0), pending(max_pos));
162  const auto& rest_indices = pending.tail(pending.size() - 1);
163  EArrX xx1 = GetSubArray(x1, rest_indices).cwiseMax(x1[i]);
164  EArrX yy1 = GetSubArray(y1, rest_indices).cwiseMax(y1[i]);
165  EArrX xx2 = GetSubArray(x2, rest_indices).cwiseMin(x2[i]);
166  EArrX yy2 = GetSubArray(y2, rest_indices).cwiseMin(y2[i]);
167 
168  EArrX w = (xx2 - xx1 + 1.0).cwiseMax(0.0);
169  EArrX h = (yy2 - yy1 + 1.0).cwiseMax(0.0);
170  EArrX inter = w * h;
171  EArrX ovr = inter / (areas[i] + GetSubArray(areas, rest_indices) - inter);
172 
173  // Update scores based on computed IoU, overlap threshold and NMS method
174  for (int j = 0; j < rest_indices.size(); ++j) {
175  typename Derived2::Scalar weight;
176  switch (method) {
177  case 1: // Linear
178  weight = (ovr(j) > overlap_thresh) ? (1.0 - ovr(j)) : 1.0;
179  break;
180  case 2: // Gaussian
181  weight = std::exp(-1.0 * ovr(j) * ovr(j) / sigma);
182  break;
183  default: // Original NMS
184  weight = (ovr(j) > overlap_thresh) ? 0.0 : 1.0;
185  }
186  (*out_scores)(rest_indices[j]) *= weight;
187  }
188 
189  // Discard boxes with new scores below min threshold and update pending
190  // indices
191  const auto& rest_scores = GetSubArray(*out_scores, rest_indices);
192  const auto& inds = GetArrayIndices(rest_scores >= score_thresh);
193  pending = GetSubArray(rest_indices, AsEArrXt(inds));
194  }
195 
196  return keep;
197 }
198 
199 template <class Derived1, class Derived2, class Derived3>
200 std::vector<int> soft_nms_cpu(
201  Eigen::ArrayBase<Derived3>* out_scores,
202  const Eigen::ArrayBase<Derived1>& proposals,
203  const Eigen::ArrayBase<Derived2>& scores,
204  float sigma = 0.5,
205  float overlap_thresh = 0.3,
206  float score_thresh = 0.001,
207  unsigned int method = 1,
208  int topN = -1) {
209  std::vector<int> indices(proposals.rows());
210  std::iota(indices.begin(), indices.end(), 0);
211  return soft_nms_cpu(
212  out_scores,
213  proposals,
214  scores,
215  indices,
216  sigma,
217  overlap_thresh,
218  score_thresh,
219  method,
220  topN);
221 }
222 
223 } // namespace utils
224 } // namespace caffe2
225 
226 #endif // CAFFE2_OPERATORS_UTILS_NMS_H_
Copyright (c) 2016-present, Facebook, Inc.