1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 #ifndef TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_NON_MAX_SUPPRESSION_H_
16 #define TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_NON_MAX_SUPPRESSION_H_
17
18 #include <algorithm>
19 #include <cmath>
20 #include <deque>
21 #include <queue>
22
23 namespace tflite {
24 namespace reference_ops {
25
26 // A pair of diagonal corners of the box.
27 struct BoxCornerEncoding {
28 float y1;
29 float x1;
30 float y2;
31 float x2;
32 };
33
ComputeIntersectionOverUnion(const float * boxes,const int i,const int j)34 inline float ComputeIntersectionOverUnion(const float* boxes, const int i,
35 const int j) {
36 auto& box_i = reinterpret_cast<const BoxCornerEncoding*>(boxes)[i];
37 auto& box_j = reinterpret_cast<const BoxCornerEncoding*>(boxes)[j];
38 const float box_i_y_min = std::min<float>(box_i.y1, box_i.y2);
39 const float box_i_y_max = std::max<float>(box_i.y1, box_i.y2);
40 const float box_i_x_min = std::min<float>(box_i.x1, box_i.x2);
41 const float box_i_x_max = std::max<float>(box_i.x1, box_i.x2);
42 const float box_j_y_min = std::min<float>(box_j.y1, box_j.y2);
43 const float box_j_y_max = std::max<float>(box_j.y1, box_j.y2);
44 const float box_j_x_min = std::min<float>(box_j.x1, box_j.x2);
45 const float box_j_x_max = std::max<float>(box_j.x1, box_j.x2);
46
47 const float area_i =
48 (box_i_y_max - box_i_y_min) * (box_i_x_max - box_i_x_min);
49 const float area_j =
50 (box_j_y_max - box_j_y_min) * (box_j_x_max - box_j_x_min);
51 if (area_i <= 0 || area_j <= 0) return 0.0;
52 const float intersection_ymax = std::min<float>(box_i_y_max, box_j_y_max);
53 const float intersection_xmax = std::min<float>(box_i_x_max, box_j_x_max);
54 const float intersection_ymin = std::max<float>(box_i_y_min, box_j_y_min);
55 const float intersection_xmin = std::max<float>(box_i_x_min, box_j_x_min);
56 const float intersection_area =
57 std::max<float>(intersection_ymax - intersection_ymin, 0.0) *
58 std::max<float>(intersection_xmax - intersection_xmin, 0.0);
59 return intersection_area / (area_i + area_j - intersection_area);
60 }
61
62 // Implements (Single-Class) Soft NMS (with Gaussian weighting).
63 // Supports functionality of TensorFlow ops NonMaxSuppressionV4 & V5.
64 // Reference: "Soft-NMS - Improving Object Detection With One Line of Code"
65 // [Bodla et al, https://arxiv.org/abs/1704.04503]
66 // Implementation adapted from the TensorFlow NMS code at
67 // tensorflow/core/kernels/non_max_suppression_op.cc.
68 //
69 // Arguments:
70 // boxes: box encodings in format [y1, x1, y2, x2], shape: [num_boxes, 4]
71 // num_boxes: number of candidates
72 // scores: scores for candidate boxes, in the same order. shape: [num_boxes]
73 // max_output_size: the maximum number of selections.
74 // iou_threshold: Intersection-over-Union (IoU) threshold for NMS
75 // score_threshold: All candidate scores below this value are rejected
76 // soft_nms_sigma: Soft NMS parameter, used for decaying scores
77 //
78 // Outputs:
79 // selected_indices: all the selected indices. Underlying array must have
80 // length >= max_output_size. Cannot be null.
81 // selected_scores: scores of selected indices. Defer from original value for
82 // Soft NMS. If not null, array must have length >= max_output_size.
83 // num_selected_indices: Number of selections. Only these many elements are
84 // set in selected_indices, selected_scores. Cannot be null.
85 //
86 // Assumes inputs are valid (for eg, iou_threshold must be >= 0).
NonMaxSuppression(const float * boxes,const int num_boxes,const float * scores,const int max_output_size,const float iou_threshold,const float score_threshold,const float soft_nms_sigma,int * selected_indices,float * selected_scores,int * num_selected_indices)87 inline void NonMaxSuppression(const float* boxes, const int num_boxes,
88 const float* scores, const int max_output_size,
89 const float iou_threshold,
90 const float score_threshold,
91 const float soft_nms_sigma, int* selected_indices,
92 float* selected_scores,
93 int* num_selected_indices) {
94 struct Candidate {
95 int index;
96 float score;
97 int suppress_begin_index;
98 };
99
100 // Priority queue to hold candidates.
101 auto cmp = [](const Candidate bs_i, const Candidate bs_j) {
102 return bs_i.score < bs_j.score;
103 };
104 std::priority_queue<Candidate, std::deque<Candidate>, decltype(cmp)>
105 candidate_priority_queue(cmp);
106 // Populate queue with candidates above the score threshold.
107 for (int i = 0; i < num_boxes; ++i) {
108 if (scores[i] > score_threshold) {
109 candidate_priority_queue.emplace(Candidate({i, scores[i], 0}));
110 }
111 }
112
113 *num_selected_indices = 0;
114 int num_outputs = std::min(static_cast<int>(candidate_priority_queue.size()),
115 max_output_size);
116 if (num_outputs == 0) return;
117
118 // NMS loop.
119 float scale = 0;
120 if (soft_nms_sigma > 0.0) {
121 scale = -0.5 / soft_nms_sigma;
122 }
123 while (*num_selected_indices < num_outputs &&
124 !candidate_priority_queue.empty()) {
125 Candidate next_candidate = candidate_priority_queue.top();
126 const float original_score = next_candidate.score;
127 candidate_priority_queue.pop();
128
129 // Overlapping boxes are likely to have similar scores, therefore we
130 // iterate through the previously selected boxes backwards in order to
131 // see if `next_candidate` should be suppressed. We also enforce a property
132 // that a candidate can be suppressed by another candidate no more than
133 // once via `suppress_begin_index` which tracks which previously selected
134 // boxes have already been compared against next_candidate prior to a given
135 // iteration. These previous selected boxes are then skipped over in the
136 // following loop.
137 bool should_hard_suppress = false;
138 for (int j = *num_selected_indices - 1;
139 j >= next_candidate.suppress_begin_index; --j) {
140 const float iou = ComputeIntersectionOverUnion(
141 boxes, next_candidate.index, selected_indices[j]);
142
143 // First decide whether to perform hard suppression.
144 if (iou >= iou_threshold) {
145 should_hard_suppress = true;
146 break;
147 }
148
149 // Suppress score if NMS sigma > 0.
150 if (soft_nms_sigma > 0.0) {
151 next_candidate.score =
152 next_candidate.score * std::exp(scale * iou * iou);
153 }
154
155 // If score has fallen below score_threshold, it won't be pushed back into
156 // the queue.
157 if (next_candidate.score <= score_threshold) break;
158 }
159 // If `next_candidate.score` has not dropped below `score_threshold`
160 // by this point, then we know that we went through all of the previous
161 // selections and can safely update `suppress_begin_index` to
162 // `selected.size()`. If on the other hand `next_candidate.score`
163 // *has* dropped below the score threshold, then since `suppress_weight`
164 // always returns values in [0, 1], further suppression by items that were
165 // not covered in the above for loop would not have caused the algorithm
166 // to select this item. We thus do the same update to
167 // `suppress_begin_index`, but really, this element will not be added back
168 // into the priority queue.
169 next_candidate.suppress_begin_index = *num_selected_indices;
170
171 if (!should_hard_suppress) {
172 if (next_candidate.score == original_score) {
173 // Suppression has not occurred, so select next_candidate.
174 selected_indices[*num_selected_indices] = next_candidate.index;
175 if (selected_scores) {
176 selected_scores[*num_selected_indices] = next_candidate.score;
177 }
178 ++*num_selected_indices;
179 }
180 if (next_candidate.score > score_threshold) {
181 // Soft suppression might have occurred and current score is still
182 // greater than score_threshold; add next_candidate back onto priority
183 // queue.
184 candidate_priority_queue.push(next_candidate);
185 }
186 }
187 }
188 }
189
190 } // namespace reference_ops
191 } // namespace tflite
192
193 #endif // TENSORFLOW_LITE_KERNELS_INTERNAL_REFERENCE_NON_MAX_SUPPRESSION_H_
194