chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,219 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user