wehub-resource-sync
488 lines
18 KiB
Python
488 lines
18 KiB
Python
# 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.
|
|
# pylint: disable=invalid-name
|
|
# ruff: noqa: E741
|
|
"""Common utilities used in Non-maximum suppression operators"""
|
|
|
|
import tvm
|
|
from tvm import te
|
|
from tvm.script.ir_builder import IRBuilder
|
|
from tvm.script.ir_builder import tirx as T
|
|
|
|
|
|
def _get_boundaries(output, box_idx):
|
|
l = tvm.te.min(
|
|
output[box_idx],
|
|
output[box_idx + 2],
|
|
)
|
|
t = tvm.te.min(
|
|
output[box_idx + 1],
|
|
output[box_idx + 3],
|
|
)
|
|
r = tvm.te.max(
|
|
output[box_idx],
|
|
output[box_idx + 2],
|
|
)
|
|
b = tvm.te.max(
|
|
output[box_idx + 1],
|
|
output[box_idx + 3],
|
|
)
|
|
return l, t, r, b
|
|
|
|
|
|
def calculate_overlap(out_tensor, box_a_idx, box_b_idx):
|
|
"""Calculate overlap of two boxes."""
|
|
a_l, a_t, a_r, a_b = _get_boundaries(out_tensor, box_a_idx)
|
|
b_l, b_t, b_r, b_b = _get_boundaries(out_tensor, box_b_idx)
|
|
|
|
# Overlapping width and height
|
|
w = tvm.te.max(0.0, tvm.te.min(a_r, b_r) - tvm.te.max(a_l, b_l))
|
|
h = tvm.te.max(0.0, tvm.te.min(a_b, b_b) - tvm.te.max(a_t, b_t))
|
|
|
|
# Overlapping area
|
|
area = h * w
|
|
|
|
# total area of the figure formed by box a and box b
|
|
# except for overlapping area
|
|
u = (a_r - a_l) * (a_b - a_t) + (b_r - b_l) * (b_b - b_t) - area
|
|
return tvm.tirx.Select(u <= 0.0, 0.0, area / u)
|
|
|
|
|
|
def binary_search(y, num_boxes, scores, score_threshold, out):
|
|
"""Binary search for score_threshold on scores sorted in descending order.
|
|
|
|
Must be called within an IRBuilder context.
|
|
"""
|
|
out = T.buffer_proxy(out)
|
|
lo_buf = T.decl_buffer([1], "int32", scope="local")
|
|
hi_buf = T.decl_buffer([1], "int32", scope="local")
|
|
lo = T.buffer_proxy(lo_buf)
|
|
hi = T.buffer_proxy(hi_buf)
|
|
lo[0] = T.int32(0)
|
|
hi[0] = tvm.tirx.Cast("int32", num_boxes)
|
|
with T.While(lo[0] < hi[0]):
|
|
mid = (hi[0] + lo[0]) >> 1
|
|
with T.If(scores[y, mid] > score_threshold):
|
|
with T.Then():
|
|
lo[0] = mid + 1
|
|
with T.Else():
|
|
hi[0] = mid
|
|
out[y] = lo[0]
|
|
|
|
|
|
def _estimate_max_detections(batch_class, input_image_size=None):
|
|
"""Estimate maximum detections based on input image size and number of classes.
|
|
|
|
This provides a more intelligent default for production environments.
|
|
"""
|
|
if input_image_size is not None:
|
|
# Estimate based on image size: larger images typically have more objects
|
|
if len(input_image_size) >= 2:
|
|
height, width = input_image_size[-2], input_image_size[-1]
|
|
total_pixels = height * width
|
|
|
|
# Base estimation per class based on image size
|
|
if total_pixels < 300000: # Small images (< 300k pixels)
|
|
base_detections_per_class = min(50, max(10, total_pixels // 2000))
|
|
elif total_pixels < 1000000: # Medium images (< 1M pixels)
|
|
base_detections_per_class = min(100, max(25, total_pixels // 3000))
|
|
else: # Large images (>= 1M pixels)
|
|
base_detections_per_class = min(200, max(50, total_pixels // 4000))
|
|
|
|
# Scale down for many classes (more realistic for multi-class scenarios)
|
|
if batch_class > 20:
|
|
# For many classes, reduce per-class detections to avoid explosion
|
|
detections_per_class = min(base_detections_per_class, 50)
|
|
else:
|
|
detections_per_class = base_detections_per_class
|
|
else:
|
|
detections_per_class = 50 # fallback
|
|
else:
|
|
# Fallback to class-based estimation
|
|
if batch_class == 1:
|
|
detections_per_class = 100 # Single class detection
|
|
elif batch_class <= 10:
|
|
detections_per_class = 50 # Small multi-class
|
|
else:
|
|
detections_per_class = 25 # Large multi-class (COCO-like)
|
|
|
|
return batch_class * detections_per_class
|
|
|
|
|
|
def collect_selected_indices(
|
|
num_class,
|
|
selected_indices,
|
|
num_detections,
|
|
row_offsets,
|
|
ir,
|
|
max_output_boxes_per_class=None,
|
|
output_shape=None,
|
|
num_total_detections=None,
|
|
input_image_size=None,
|
|
):
|
|
"""Collect selected indices from the core NMS loop into one linear output.
|
|
|
|
Parameters
|
|
----------
|
|
num_class : int
|
|
selected_indices : tvm.te.Tensor
|
|
2-D tensor with shape (batch_size * num_classes, num_boxes), representing the indices
|
|
of selected boxes by the core NMS loop.
|
|
num_detections : tvm.te.Tensor
|
|
1-D tensor with shape (batch_size * num_classes,), representing
|
|
the number of boxes selected by the core NMS loop, per batch and class.
|
|
row_offsets : tvm.te.Tensor
|
|
1-D tensor with shape (batch_size * num_classes,), this should be the exclusive scan
|
|
of num_detections.
|
|
ir : function
|
|
A function to generate IR for CPU or GPU, see its usage in vision/nms.py and cuda/nms.py.
|
|
|
|
Returns
|
|
-------
|
|
out : tvm.te.Tensor
|
|
The output is indices of size (batch_size * num_class* num_boxes , 3).
|
|
Rows of indices are ordered such that selected boxes from batch 0, class 0 come
|
|
first, in descending of scores, followed by boxes from batch 0, class 1 etc.
|
|
"""
|
|
batch_class, num_boxes = selected_indices.shape
|
|
|
|
if output_shape is not None:
|
|
return te.extern(
|
|
[output_shape],
|
|
[selected_indices, num_detections, row_offsets],
|
|
lambda ins, outs: ir(
|
|
num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
|
|
),
|
|
dtype=["int64"],
|
|
name="collect_indices",
|
|
tag="collect_indices",
|
|
)
|
|
|
|
# TODO: Implement dynamic trimming based on num_total_detections
|
|
if num_total_detections is not None:
|
|
if isinstance(max_output_boxes_per_class, int):
|
|
out_rows = batch_class * max_output_boxes_per_class
|
|
else:
|
|
# Smart fallback based on input image size and typical production scenarios
|
|
out_rows = _estimate_max_detections(batch_class, input_image_size)
|
|
|
|
return te.extern(
|
|
[(out_rows, 3)],
|
|
[selected_indices, num_detections, row_offsets],
|
|
lambda ins, outs: ir(
|
|
num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
|
|
),
|
|
dtype=["int64"],
|
|
name="collect_indices",
|
|
tag="collect_indices",
|
|
)
|
|
|
|
if isinstance(max_output_boxes_per_class, int):
|
|
out_rows = batch_class * max_output_boxes_per_class
|
|
return te.extern(
|
|
[(out_rows, 3)],
|
|
[selected_indices, num_detections, row_offsets],
|
|
lambda ins, outs: ir(
|
|
num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
|
|
),
|
|
dtype=["int64"],
|
|
name="collect_indices",
|
|
tag="collect_indices",
|
|
)
|
|
|
|
if isinstance(max_output_boxes_per_class, te.Tensor):
|
|
try:
|
|
if len(max_output_boxes_per_class.shape) == 0:
|
|
max_boxes_val = int(max_output_boxes_per_class.data.numpy())
|
|
elif (
|
|
len(max_output_boxes_per_class.shape) == 1
|
|
and max_output_boxes_per_class.shape[0] == 1
|
|
):
|
|
max_boxes_val = int(max_output_boxes_per_class.data.numpy()[0])
|
|
else:
|
|
max_boxes_val = num_boxes
|
|
except (ValueError, IndexError, AttributeError):
|
|
max_boxes_val = num_boxes
|
|
|
|
out_rows = batch_class * max_boxes_val
|
|
return te.extern(
|
|
[(out_rows, 3)],
|
|
[selected_indices, num_detections, row_offsets],
|
|
lambda ins, outs: ir(
|
|
num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
|
|
),
|
|
dtype=["int64"],
|
|
name="collect_indices",
|
|
tag="collect_indices",
|
|
)
|
|
|
|
return te.extern(
|
|
[(batch_class * num_boxes, 3)],
|
|
[selected_indices, num_detections, row_offsets],
|
|
lambda ins, outs: ir(
|
|
num_class, ins[0], ins[1], ins[2], outs[0], max_output_boxes_per_class
|
|
),
|
|
dtype=["int64"],
|
|
name="collect_indices",
|
|
tag="collect_indices",
|
|
)
|
|
|
|
|
|
def collect_selected_indices_and_scores(
|
|
selected_indices, selected_scores, num_detections, row_offsets, num_total_detections, ir
|
|
):
|
|
"""Collect selected indices and scores from the core NMS loop into one linear output.
|
|
|
|
Parameters
|
|
----------
|
|
selected_indices : tvm.te.Tensor
|
|
2-D tensor with shape (batch_size * num_classes, num_boxes), representing the indices
|
|
of selected boxes by the core NMS loop.
|
|
selected_scores : tvm.te.Tensor
|
|
2-D tensor with shape (batch_size * num_classes, num_boxes), representing the scores
|
|
of selected boxes by the core NMS loop.
|
|
num_detections : tvm.te.Tensor
|
|
2-D tensor with shape (batch_size, num_classes), representing
|
|
the number of boxes selected by the core NMS loop, per batch and class.
|
|
row_offsets : tvm.te.Tensor
|
|
2-D tensor with shape (batch_size, num_classes), this should be the exclusive scan
|
|
of num_detections along axis 1.
|
|
num_total_detections : tvm.te.Tensor
|
|
Total number of detections.
|
|
ir : function
|
|
A function to generate IR for CPU or GPU, see its usage in vision/nms.py and cuda/nms.py.
|
|
|
|
Returns
|
|
-------
|
|
out : [tvm.te.Tensor, tvm.te.Tensor]
|
|
The output is two tensors. The first is indices of size
|
|
(batch_size, num_class* num_boxes, 2), and the second is scores of size
|
|
(batch_size, num_class* num_boxes).
|
|
"""
|
|
batch_size, num_class = row_offsets.shape
|
|
num_boxes = selected_indices.shape[1]
|
|
return te.extern(
|
|
[(batch_size, num_class * num_boxes, 2), (batch_size, num_class * num_boxes)],
|
|
[selected_indices, selected_scores, num_detections, row_offsets, num_total_detections],
|
|
lambda ins, outs: ir(ins[0], ins[1], ins[2], ins[3], ins[4], outs[0], outs[1]),
|
|
dtype=["int64", "float32"],
|
|
name="collect_indices_and_scores",
|
|
tag="collect_indices_and_scores",
|
|
)
|
|
|
|
|
|
def _all_class_nms_ir(
|
|
boxes,
|
|
sorted_scores,
|
|
sorted_indices,
|
|
valid_count,
|
|
batch_class,
|
|
num_class,
|
|
num_anchors,
|
|
iou_threshold,
|
|
max_output_size_per_class,
|
|
box_indices,
|
|
selected_scores,
|
|
num_valid_boxes,
|
|
nms_loop,
|
|
score_threshold=None,
|
|
):
|
|
with IRBuilder() as ib:
|
|
# Wrap buffers with T.buffer_proxy for flat indexing support
|
|
boxes = T.buffer_proxy(boxes)
|
|
box_indices = T.buffer_proxy(box_indices)
|
|
if selected_scores is not None:
|
|
selected_scores = T.buffer_proxy(selected_scores)
|
|
|
|
if isinstance(iou_threshold, float | int):
|
|
iou_threshold = tvm.tirx.FloatImm("float32", float(iou_threshold))
|
|
elif isinstance(iou_threshold, te.Tensor):
|
|
if len(iou_threshold.shape) == 0:
|
|
iou_threshold = iou_threshold()
|
|
elif len(iou_threshold.shape) == 1 and iou_threshold.shape[0] == 1:
|
|
iou_threshold = iou_threshold[0]
|
|
else:
|
|
iou_threshold = tvm.tirx.FloatImm("float32", 0.5)
|
|
|
|
if isinstance(max_output_size_per_class, int):
|
|
max_output_size_per_class = tvm.tirx.const(max_output_size_per_class)
|
|
elif isinstance(max_output_size_per_class, te.Tensor):
|
|
if len(max_output_size_per_class.shape) == 0:
|
|
max_output_size_per_class = max_output_size_per_class()
|
|
elif (
|
|
len(max_output_size_per_class.shape) == 1
|
|
and max_output_size_per_class.shape[0] == 1
|
|
):
|
|
max_output_size_per_class = max_output_size_per_class[0]
|
|
else:
|
|
max_output_size_per_class = tvm.tirx.const(1000)
|
|
|
|
def calc_overlap(i, j, k):
|
|
offset_j = sorted_indices[i, j] * 4
|
|
offset_k = sorted_indices[i, k] * 4
|
|
batch_id = i // num_class
|
|
base_bbox_idx = batch_id * num_anchors * 4
|
|
return calculate_overlap(
|
|
boxes,
|
|
base_bbox_idx + offset_j,
|
|
base_bbox_idx + offset_k,
|
|
)
|
|
|
|
def on_new_valid_box(tid, num_current_valid_box, i, j):
|
|
with T.If(tid + 0 == 0):
|
|
with T.Then():
|
|
box_indices[i, num_current_valid_box] = sorted_indices[i, j]
|
|
if selected_scores is not None:
|
|
selected_scores[i, num_current_valid_box] = sorted_scores[i, j]
|
|
|
|
def on_new_invalidated_box(*_):
|
|
pass
|
|
|
|
def needs_bbox_check(*_):
|
|
return tvm.tirx.const(True)
|
|
|
|
nms_loop(
|
|
batch_class,
|
|
tvm.tirx.IntImm("int32", -1), # top_k
|
|
iou_threshold,
|
|
max_output_size_per_class,
|
|
valid_count,
|
|
on_new_valid_box,
|
|
on_new_invalidated_box,
|
|
needs_bbox_check,
|
|
calc_overlap,
|
|
sorted_scores,
|
|
num_valid_boxes,
|
|
score_threshold,
|
|
)
|
|
|
|
return ib.get()
|
|
|
|
|
|
def run_all_class_nms(
|
|
boxes,
|
|
sorted_scores,
|
|
sorted_indices,
|
|
valid_count,
|
|
max_output_size_per_class,
|
|
iou_threshold,
|
|
nms_loop,
|
|
return_scores=False,
|
|
score_threshold=None,
|
|
):
|
|
"""The core all class NMS routine.
|
|
|
|
Parameters
|
|
----------
|
|
boxes : tvm.te.Tensor
|
|
3-D tensor with shape (batch_size, num_boxes, 4)
|
|
sorted_scores : tvm.te.Tensor
|
|
2-D tensor with shape (batch_size * num_classes, num_boxes).
|
|
One of the outputs from argsort.
|
|
sorted_indices : tvm.te.Tensor
|
|
2-D tensor with shape (batch_size * num_classes, num_boxes).
|
|
The other output from argsort.
|
|
valid_count : tvm.te.Tensor
|
|
1-D tensor with shape (batch_size * num_classes,), representing
|
|
the number of boxes whose score is above score_threshold, per batch and class.
|
|
max_output_size_per_class : int or tvm.te.Tensor, optional
|
|
The maxinum number of output selected boxes per class.
|
|
iou_threshold : float or tvm.te.Tensor, optional
|
|
IoU test threshold.
|
|
nms_loop : function
|
|
A core NMS loop, see its usage in vision/nms.py and cuda/nms.py.
|
|
return_scores : bool, optional
|
|
Whether or not to return selected scores, needed by the tensorflow output format.
|
|
|
|
Returns
|
|
-------
|
|
out : a list of tvm.te.Tensor
|
|
The output is three tensors, the first and second are indices and scores of size
|
|
(batch_size * num_class, num_boxes), and the third is a tensor
|
|
num_selected_boxes of shape (batch_size * num_class,) representing the total number of
|
|
selected boxes per batch and class. If return_scores is False, the second output is
|
|
None.
|
|
"""
|
|
batch, num_boxes, _ = boxes.shape
|
|
batch_class = sorted_scores.shape[0]
|
|
num_class = batch_class // batch
|
|
|
|
if return_scores is False:
|
|
all_class_num0_buf = tvm.tirx.decl_buffer(
|
|
(batch_class, num_boxes), "int32", "all_class_nms0", data_alignment=8, layout=None
|
|
)
|
|
all_class_num1_buf = tvm.tirx.decl_buffer(
|
|
(batch_class,), "int32", "all_class_nms1", data_alignment=8, layout=None
|
|
)
|
|
extern_inputs = [boxes, sorted_scores, sorted_indices, valid_count]
|
|
if score_threshold is not None:
|
|
extern_inputs.append(score_threshold)
|
|
|
|
selected_indices, num_detections = te.extern(
|
|
[(batch_class, num_boxes), (batch_class,)],
|
|
extern_inputs,
|
|
lambda ins, outs: _all_class_nms_ir(
|
|
ins[0], # boxes
|
|
ins[1], # sorted_scores
|
|
ins[2], # sorted_indices
|
|
ins[3], # valid_count
|
|
batch_class,
|
|
num_class,
|
|
num_boxes,
|
|
iou_threshold,
|
|
max_output_size_per_class,
|
|
outs[0], # box_indices
|
|
None, # scores
|
|
outs[1], # num_selected_boxes
|
|
nms_loop,
|
|
ins[4] if score_threshold is not None else None, # score_threshold
|
|
),
|
|
out_buffers=[all_class_num0_buf, all_class_num1_buf],
|
|
dtype=["int32", "int32"],
|
|
name="all_class_nms",
|
|
tag="all_class_nms",
|
|
)
|
|
return selected_indices, None, num_detections
|
|
|
|
extern_inputs = [boxes, sorted_scores, sorted_indices, valid_count]
|
|
if score_threshold is not None:
|
|
extern_inputs.append(score_threshold)
|
|
|
|
return te.extern(
|
|
[(batch_class, num_boxes), (batch_class, num_boxes), (batch_class,)],
|
|
extern_inputs,
|
|
lambda ins, outs: _all_class_nms_ir(
|
|
ins[0], # boxes
|
|
ins[1], # sorted_scores
|
|
ins[2], # sorted_indices
|
|
ins[3], # valid_count
|
|
batch_class,
|
|
num_class,
|
|
num_boxes,
|
|
iou_threshold,
|
|
max_output_size_per_class,
|
|
outs[0], # box_indices
|
|
outs[1], # selected scores
|
|
outs[2], # num_selected_boxes
|
|
nms_loop,
|
|
ins[4] if score_threshold is not None else None, # score_threshold
|
|
),
|
|
dtype=["int32", "float32", "int32"],
|
|
name="all_class_nms",
|
|
tag="all_class_nms",
|
|
)
|