mirror of https://github.com/open-mmlab/mmcv.git
206 lines
7.6 KiB
Python
206 lines
7.6 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from mmcv.utils import IS_CUDA_AVAILABLE, IS_MLU_AVAILABLE
|
|
|
|
|
|
class Testnms:
|
|
|
|
@pytest.mark.parametrize('device', [
|
|
pytest.param(
|
|
'cuda',
|
|
marks=pytest.mark.skipif(
|
|
not IS_CUDA_AVAILABLE, reason='requires CUDA support')),
|
|
pytest.param(
|
|
'mlu',
|
|
marks=pytest.mark.skipif(
|
|
not IS_MLU_AVAILABLE, reason='requires MLU support'))
|
|
])
|
|
def test_nms_allclose(self, device):
|
|
from mmcv.ops import nms
|
|
np_boxes = np.array([[6.0, 3.0, 8.0, 7.0], [3.0, 6.0, 9.0, 11.0],
|
|
[3.0, 7.0, 10.0, 12.0], [1.0, 4.0, 13.0, 7.0]],
|
|
dtype=np.float32)
|
|
np_scores = np.array([0.6, 0.9, 0.7, 0.2], dtype=np.float32)
|
|
np_inds = np.array([1, 0, 3])
|
|
np_dets = np.array([[3.0, 6.0, 9.0, 11.0, 0.9],
|
|
[6.0, 3.0, 8.0, 7.0, 0.6],
|
|
[1.0, 4.0, 13.0, 7.0, 0.2]])
|
|
boxes = torch.from_numpy(np_boxes)
|
|
scores = torch.from_numpy(np_scores)
|
|
dets, inds = nms(boxes, scores, iou_threshold=0.3, offset=0)
|
|
assert np.allclose(dets, np_dets) # test cpu
|
|
assert np.allclose(inds, np_inds) # test cpu
|
|
dets, inds = nms(
|
|
boxes.to(device), scores.to(device), iou_threshold=0.3, offset=0)
|
|
assert np.allclose(dets.cpu().numpy(), np_dets) # test gpu
|
|
assert np.allclose(inds.cpu().numpy(), np_inds) # test gpu
|
|
|
|
def test_softnms_allclose(self):
|
|
if not torch.cuda.is_available():
|
|
return
|
|
from mmcv.ops import soft_nms
|
|
np_boxes = np.array([[6.0, 3.0, 8.0, 7.0], [3.0, 6.0, 9.0, 11.0],
|
|
[3.0, 7.0, 10.0, 12.0], [1.0, 4.0, 13.0, 7.0]],
|
|
dtype=np.float32)
|
|
np_scores = np.array([0.6, 0.9, 0.7, 0.2], dtype=np.float32)
|
|
|
|
np_output = {
|
|
'linear': {
|
|
'dets':
|
|
np.array(
|
|
[[3., 6., 9., 11., 0.9], [6., 3., 8., 7., 0.6],
|
|
[3., 7., 10., 12., 0.29024392], [1., 4., 13., 7., 0.2]],
|
|
dtype=np.float32),
|
|
'inds':
|
|
np.array([1, 0, 2, 3], dtype=np.int64)
|
|
},
|
|
'gaussian': {
|
|
'dets':
|
|
np.array([[3., 6., 9., 11., 0.9], [6., 3., 8., 7., 0.59630775],
|
|
[3., 7., 10., 12., 0.35275510],
|
|
[1., 4., 13., 7., 0.18650459]],
|
|
dtype=np.float32),
|
|
'inds':
|
|
np.array([1, 0, 2, 3], dtype=np.int64)
|
|
},
|
|
'naive': {
|
|
'dets':
|
|
np.array([[3., 6., 9., 11., 0.9], [6., 3., 8., 7., 0.6],
|
|
[1., 4., 13., 7., 0.2]],
|
|
dtype=np.float32),
|
|
'inds':
|
|
np.array([1, 0, 3], dtype=np.int64)
|
|
}
|
|
}
|
|
|
|
boxes = torch.from_numpy(np_boxes)
|
|
scores = torch.from_numpy(np_scores)
|
|
|
|
configs = [[0.3, 0.5, 0.01, 'linear'], [0.3, 0.5, 0.01, 'gaussian'],
|
|
[0.3, 0.5, 0.01, 'naive']]
|
|
|
|
for iou, sig, mscore, m in configs:
|
|
dets, inds = soft_nms(
|
|
boxes,
|
|
scores,
|
|
iou_threshold=iou,
|
|
sigma=sig,
|
|
min_score=mscore,
|
|
method=m)
|
|
assert np.allclose(dets.cpu().numpy(), np_output[m]['dets'])
|
|
assert np.allclose(inds.cpu().numpy(), np_output[m]['inds'])
|
|
|
|
if torch.__version__ != 'parrots':
|
|
boxes = boxes.cuda()
|
|
scores = scores.cuda()
|
|
for iou, sig, mscore, m in configs:
|
|
dets, inds = soft_nms(
|
|
boxes,
|
|
scores,
|
|
iou_threshold=iou,
|
|
sigma=sig,
|
|
min_score=mscore,
|
|
method=m)
|
|
assert np.allclose(dets.cpu().numpy(), np_output[m]['dets'])
|
|
assert np.allclose(inds.cpu().numpy(), np_output[m]['inds'])
|
|
|
|
def test_nms_match(self):
|
|
if not torch.cuda.is_available():
|
|
return
|
|
from mmcv.ops import nms, nms_match
|
|
iou_thr = 0.6
|
|
# empty input
|
|
empty_dets = np.array([])
|
|
assert len(nms_match(empty_dets, iou_thr)) == 0
|
|
|
|
# non empty ndarray input
|
|
np_dets = np.array(
|
|
[[49.1, 32.4, 51.0, 35.9, 0.9], [49.3, 32.9, 51.0, 35.3, 0.9],
|
|
[35.3, 11.5, 39.9, 14.5, 0.4], [35.2, 11.7, 39.7, 15.7, 0.3]],
|
|
dtype=np.float32)
|
|
np_groups = nms_match(np_dets, iou_thr)
|
|
assert isinstance(np_groups[0], np.ndarray)
|
|
assert len(np_groups) == 2
|
|
tensor_dets = torch.from_numpy(np_dets)
|
|
boxes = tensor_dets[:, :4]
|
|
scores = tensor_dets[:, 4]
|
|
nms_keep_inds = nms(boxes.contiguous(), scores.contiguous(),
|
|
iou_thr)[1]
|
|
assert {g[0].item() for g in np_groups} == set(nms_keep_inds.tolist())
|
|
|
|
# non empty tensor input
|
|
tensor_dets = torch.from_numpy(np_dets)
|
|
tensor_groups = nms_match(tensor_dets, iou_thr)
|
|
assert isinstance(tensor_groups[0], torch.Tensor)
|
|
for i in range(len(tensor_groups)):
|
|
assert np.equal(tensor_groups[i].numpy(), np_groups[i]).all()
|
|
|
|
# input of wrong shape
|
|
wrong_dets = np.zeros((2, 3))
|
|
with pytest.raises(AssertionError):
|
|
nms_match(wrong_dets, iou_thr)
|
|
|
|
def test_batched_nms(self):
|
|
import mmcv
|
|
from mmcv.ops import batched_nms
|
|
results = mmcv.load('./tests/data/batched_nms_data.pkl')
|
|
|
|
nms_max_num = 100
|
|
nms_cfg = dict(
|
|
type='nms',
|
|
iou_threshold=0.7,
|
|
score_threshold=0.5,
|
|
max_num=nms_max_num)
|
|
boxes, keep = batched_nms(
|
|
torch.from_numpy(results['boxes']),
|
|
torch.from_numpy(results['scores']),
|
|
torch.from_numpy(results['idxs']),
|
|
nms_cfg,
|
|
class_agnostic=False)
|
|
|
|
nms_cfg.update(split_thr=100)
|
|
seq_boxes, seq_keep = batched_nms(
|
|
torch.from_numpy(results['boxes']),
|
|
torch.from_numpy(results['scores']),
|
|
torch.from_numpy(results['idxs']),
|
|
nms_cfg,
|
|
class_agnostic=False)
|
|
|
|
assert torch.equal(keep, seq_keep)
|
|
assert torch.equal(boxes, seq_boxes)
|
|
assert torch.equal(keep,
|
|
torch.from_numpy(results['keep'][:nms_max_num]))
|
|
|
|
nms_cfg = dict(type='soft_nms', iou_threshold=0.7)
|
|
boxes, keep = batched_nms(
|
|
torch.from_numpy(results['boxes']),
|
|
torch.from_numpy(results['scores']),
|
|
torch.from_numpy(results['idxs']),
|
|
nms_cfg,
|
|
class_agnostic=False)
|
|
|
|
nms_cfg.update(split_thr=100)
|
|
seq_boxes, seq_keep = batched_nms(
|
|
torch.from_numpy(results['boxes']),
|
|
torch.from_numpy(results['scores']),
|
|
torch.from_numpy(results['idxs']),
|
|
nms_cfg,
|
|
class_agnostic=False)
|
|
|
|
assert torch.equal(keep, seq_keep)
|
|
assert torch.equal(boxes, seq_boxes)
|
|
|
|
# test skip nms when `nms_cfg` is None
|
|
seq_boxes, seq_keep = batched_nms(
|
|
torch.from_numpy(results['boxes']),
|
|
torch.from_numpy(results['scores']),
|
|
torch.from_numpy(results['idxs']),
|
|
None,
|
|
class_agnostic=False)
|
|
assert len(seq_keep) == len(results['boxes'])
|
|
# assert score is descending order
|
|
assert ((seq_boxes[:, -1][1:] - seq_boxes[:, -1][:-1]) < 0).all()
|