# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you under the Apache License, Version 2.0 (the # "License"); you may not use this file except in compliance # with the License. You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, # software distributed under the License is distributed on an # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. """Numpy reference implementation for classic non_max_suppression.""" import numpy as np def _iou(box_a, box_b, coord_start): """Compute IoU between two boxes.""" a = box_a[coord_start : coord_start + 4] b = box_b[coord_start : coord_start + 4] a_l, a_t, a_r, a_b = min(a[0], a[2]), min(a[1], a[3]), max(a[0], a[2]), max(a[1], a[3]) b_l, b_t, b_r, b_b = min(b[0], b[2]), min(b[1], b[3]), max(b[0], b[2]), max(b[1], b[3]) w = max(0.0, min(a_r, b_r) - max(a_l, b_l)) h = max(0.0, min(a_b, b_b) - max(a_t, b_t)) area = w * h u = (a_r - a_l) * (a_b - a_t) + (b_r - b_l) * (b_b - b_t) - area return 0.0 if u <= 0 else area / u def non_max_suppression_python( data, valid_count, indices, max_output_size=-1, iou_threshold=0.5, force_suppress=False, top_k=-1, coord_start=2, score_index=1, id_index=0, return_indices=True, invalid_to_bottom=False, soft_nms_sigma=0.0, score_threshold=0.0, ): """Numpy reference for classic non_max_suppression. Parameters ---------- data : numpy.ndarray 3-D array, shape [batch_size, num_anchors, elem_length]. valid_count : numpy.ndarray 1-D array, shape [batch_size]. indices : numpy.ndarray 2-D array, shape [batch_size, num_anchors]. Returns ------- If return_indices is True and soft_nms_sigma == 0.0: (box_indices, valid_box_count) If return_indices is True and soft_nms_sigma > 0.0: (out_data, box_indices, valid_box_count) Otherwise: modified data tensor """ batch_size, num_anchors, _ = data.shape out_data = np.full_like(data, -1.0) out_box_indices = np.full((batch_size, num_anchors), -1, dtype="int32") compacted = np.full((batch_size, num_anchors), -1, dtype="int32") valid_box_count = np.zeros((batch_size, 1), dtype="int32") is_soft_nms = soft_nms_sigma > 0.0 thresh = score_threshold if is_soft_nms else 0.0 soft_nms_scale = -0.5 / soft_nms_sigma if is_soft_nms else 0.0 for i in range(batch_size): nkeep = int(valid_count[i]) if 0 < top_k < nkeep: nkeep = top_k # Sort by score descending scores = data[i, :nkeep, score_index].copy() sorted_idx = np.argsort(-scores) # Copy sorted boxes for j in range(nkeep): src = sorted_idx[j] out_data[i, j, :] = data[i, src, :] out_box_indices[i, j] = src if is_soft_nms: num_selected = 0 while num_selected < nkeep and (max_output_size < 0 or num_selected < max_output_size): best_idx = -1 best_score = thresh for j in range(num_selected, nkeep): if out_box_indices[i, j] >= 0 and out_data[i, j, score_index] > best_score: best_idx = j best_score = out_data[i, j, score_index] if best_idx < 0: break if best_idx != num_selected: out_data[i, [num_selected, best_idx], :] = out_data[ i, [best_idx, num_selected], : ] out_box_indices[i, [num_selected, best_idx]] = out_box_indices[ i, [best_idx, num_selected] ] selected_idx = num_selected for j in range(selected_idx + 1, nkeep): if out_box_indices[i, j] < 0 or out_data[i, j, score_index] <= thresh: continue do_suppress = False if force_suppress: do_suppress = True elif id_index >= 0: do_suppress = ( out_data[i, selected_idx, id_index] == out_data[i, j, id_index] ) else: do_suppress = True if not do_suppress: continue iou = _iou(out_data[i, selected_idx], out_data[i, j], coord_start) if iou >= iou_threshold: out_box_indices[i, j] = -1 else: out_data[i, j, score_index] *= np.exp(soft_nms_scale * (iou**2)) if out_data[i, j, score_index] <= thresh: out_box_indices[i, j] = -1 num_selected += 1 valid_box_count[i, 0] = num_selected if return_indices: for j in range(num_selected): orig_idx = out_box_indices[i, j] compacted[i, j] = int(indices[i, orig_idx]) out_box_indices[i, j] = compacted[i, j] for j in range(num_selected, num_anchors): out_data[i, j, :] = -1.0 out_box_indices[i, j] = -1 else: out_data[i, num_selected:, :] = -1.0 continue # Greedy NMS num_valid = 0 for j in range(nkeep): if out_data[i, j, score_index] <= thresh: out_data[i, j, :] = -1.0 out_box_indices[i, j] = -1 continue if 0 < max_output_size <= num_valid: out_data[i, j, :] = -1.0 out_box_indices[i, j] = -1 continue num_valid += 1 # Suppress overlapping boxes for k in range(j + 1, nkeep): if out_data[i, k, score_index] <= thresh: continue do_suppress = False if force_suppress: do_suppress = True elif id_index >= 0: do_suppress = out_data[i, j, id_index] == out_data[i, k, id_index] else: do_suppress = True if do_suppress: iou = _iou(out_data[i, j], out_data[i, k], coord_start) if iou >= iou_threshold: out_data[i, k, score_index] = -1.0 out_box_indices[i, k] = -1 if return_indices: # Compact valid indices to top and remap to original cnt = 0 for j in range(num_anchors): if out_box_indices[i, j] >= 0: orig_idx = out_box_indices[i, j] compacted[i, cnt] = int(indices[i, orig_idx]) cnt += 1 valid_box_count[i, 0] = cnt if return_indices: if is_soft_nms: return [out_data, compacted, valid_box_count] return [compacted, valid_box_count] if invalid_to_bottom: # Rearrange valid boxes to top result = np.full_like(data, -1.0) for i in range(batch_size): cnt = 0 for j in range(num_anchors): if out_data[i, j, score_index] >= 0: result[i, cnt, :] = out_data[i, j, :] cnt += 1 return result return out_data