75 lines
3.9 KiB
Python
75 lines
3.9 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
|
|
def test_aligned_3d_nms():
|
|
from mmdet3d.core.post_processing import aligned_3d_nms
|
|
|
|
boxes = torch.tensor([[1.2261, 0.6679, -1.2678, 2.6547, 1.0428, 0.1000],
|
|
[5.0919, 0.6512, 0.7238, 5.4821, 1.2451, 2.1095],
|
|
[6.8392, -1.2205, 0.8570, 7.6920, 0.3220, 3.2223],
|
|
[3.6900, -0.4235, -1.0380, 4.4415, 0.2671, -0.1442],
|
|
[4.8071, -1.4311, 0.7004, 5.5788, -0.6837, 1.2487],
|
|
[2.1807, -1.5811, -1.1289, 3.0151, -0.1346, -0.5351],
|
|
[4.4631, -4.2588, -1.1403, 5.3012, -3.4463, -0.3212],
|
|
[4.7607, -3.3311, 0.5993, 5.2976, -2.7874, 1.2273],
|
|
[3.1265, 0.7113, -0.0296, 3.8944, 1.3532, 0.9785],
|
|
[5.5828, -3.5350, 1.0105, 8.2841, -0.0405, 3.3614],
|
|
[3.0003, -2.1099, -1.0608, 5.3423, 0.0328, 0.6252],
|
|
[2.7148, 0.6082, -1.1738, 3.6995, 1.2375, -0.0209],
|
|
[4.9263, -0.2152, 0.2889, 5.6963, 0.3416, 1.3471],
|
|
[5.0713, 1.3459, -0.2598, 5.6278, 1.9300, 1.2835],
|
|
[4.5985, -2.3996, -0.3393, 5.2705, -1.7306, 0.5698],
|
|
[4.1386, 0.5658, 0.0422, 4.8937, 1.1983, 0.9911],
|
|
[2.7694, -1.9822, -1.0637, 4.0691, 0.3575, -0.1393],
|
|
[4.6464, -3.0123, -1.0694, 5.1421, -2.4450, -0.3758],
|
|
[3.4754, 0.4443, -1.1282, 4.6727, 1.3786, 0.2550],
|
|
[2.5905, -0.3504, -1.1202, 3.1599, 0.1153, -0.3036],
|
|
[4.1336, -3.4813, 1.1477, 6.2091, -0.8776, 2.6757],
|
|
[3.9966, 0.2069, -1.1148, 5.0841, 1.0525, -0.0648],
|
|
[4.3216, -1.8647, 0.4733, 6.2069, 0.6671, 3.3363],
|
|
[4.7683, 0.4286, -0.0500, 5.5642, 1.2906, 0.8902],
|
|
[1.7337, 0.7625, -1.0058, 3.0675, 1.3617, 0.3849],
|
|
[4.7193, -3.3687, -0.9635, 5.1633, -2.7656, 1.1001],
|
|
[4.4704, -2.7744, -1.1127, 5.0971, -2.0228, -0.3150],
|
|
[2.7027, 0.6122, -0.9169, 3.3083, 1.2117, 0.6129],
|
|
[4.8789, -2.0025, 0.8385, 5.5214, -1.3668, 1.3552],
|
|
[3.7856, -1.7582, -0.1738, 5.3373, -0.6300, 0.5558]])
|
|
|
|
scores = torch.tensor([
|
|
3.6414e-03, 2.2901e-02, 2.7576e-04, 1.2238e-02, 5.9310e-04, 1.2659e-01,
|
|
2.4104e-02, 5.0742e-03, 2.3581e-03, 2.0946e-07, 8.8039e-01, 1.9127e-01,
|
|
5.0469e-05, 9.3638e-03, 3.0663e-03, 9.4350e-03, 5.3380e-02, 1.7895e-01,
|
|
2.0048e-01, 1.1294e-03, 3.0304e-08, 2.0237e-01, 1.0894e-08, 6.7972e-02,
|
|
6.7156e-01, 9.3986e-04, 7.9470e-01, 3.9736e-01, 1.8000e-04, 7.9151e-04
|
|
])
|
|
|
|
cls = torch.tensor([
|
|
8, 8, 8, 3, 3, 1, 3, 3, 7, 8, 0, 6, 7, 8, 3, 7, 2, 7, 6, 3, 8, 6, 6, 7,
|
|
6, 8, 7, 6, 3, 1
|
|
])
|
|
|
|
pick = aligned_3d_nms(boxes, scores, cls, 0.25)
|
|
expected_pick = torch.tensor([
|
|
10, 26, 24, 27, 21, 18, 17, 5, 23, 16, 6, 1, 3, 15, 13, 7, 0, 14, 8,
|
|
19, 25, 29, 4, 2, 28, 12, 9, 20, 22
|
|
])
|
|
|
|
assert torch.all(pick == expected_pick)
|
|
|
|
|
|
def test_circle_nms():
|
|
from mmdet3d.core.post_processing import circle_nms
|
|
boxes = torch.tensor([[-11.1100, 2.1300, 0.8823],
|
|
[-11.2810, 2.2422, 0.8914],
|
|
[-10.3966, -0.3198, 0.8643],
|
|
[-10.2906, -13.3159,
|
|
0.8401], [5.6518, 9.9791, 0.8271],
|
|
[-11.2652, 13.3637, 0.8267],
|
|
[4.7768, -13.0409, 0.7810], [5.6621, 9.0422, 0.7753],
|
|
[-10.5561, 18.9627, 0.7518],
|
|
[-10.5643, 13.2293, 0.7200]])
|
|
keep = circle_nms(boxes.numpy(), 0.175)
|
|
expected_keep = [1, 2, 3, 4, 5, 6, 7, 8, 9]
|
|
assert np.all(keep == expected_keep)
|