def batched_nms(boxes: Tensor,
scores: Tensor,
idxs: Tensor,
iou_thresh: float,
max_output: Optional[int] = None) -> Tensor:
r"""
Performs non-maximum suppression (NMS) on the boxes according to
their intersection-over-union (IoU).
:param boxes: tensor of shape `(N, 4)`; the boxes to perform nms on;
each box is expected to be in `(x1, y1, x2, y2)` format.
:param iou_thresh: ``IoU`` threshold for overlapping.
:param idxs: tensor of shape `(N,)`, the class indexs of boxes in the batch.
:param scores: tensor of shape `(N,)`, the score of boxes.
:return: indices of the elements that have been kept by NMS.
Examples:
.. testcode::
import numpy as np
from megengine import tensor
x = np.zeros((100,4))
np.random.seed(42)
x[:,:2] = np.random.rand(100,2) * 20
x[:,2:] = np.random.rand(100,2) * 20 + 100
scores = tensor(np.random.rand(100))
idxs = tensor(np.random.randint(0, 10, 100))
inp = tensor(x)
result = batched_nms(inp, scores, idxs, iou_thresh=0.6)
print(result.numpy())
Outputs:
.. testoutput::
[75 41 99 98 69 64 11 27 35 18]
"""
assert (boxes.ndim == 2
and boxes.shape[1] == 4), "the expected shape of boxes is (N, 4)"
assert scores.ndim == 1, "the expected shape of scores is (N,)"
assert idxs.ndim == 1, "the expected shape of idxs is (N,)"
assert (boxes.shape[0] == scores.shape[0] ==
idxs.shape[0]), "number of boxes, scores and idxs are not matched"
idxs = idxs.detach()
max_coordinate = boxes.max()
offsets = idxs.astype("float32") * (max_coordinate + 1)
boxes = boxes + offsets.reshape(-1, 1)
return F.nn.nms(boxes, scores, iou_thresh, max_output)
def normalize_image(tensor: megengine.Tensor, scale=255):
"""normalize image tensors of any range to [0, scale=255]"""
mi = tensor.min()
ma = tensor.max()
tensor = scale * (tensor - mi) / (ma - mi + 1e-9)
return tensor
def test_max():
a = Tensor([1, 2])
a = a.reshape((1, 2))
assert a.max().ndim == 0
assert a.max(axis=1).ndim == 1
