Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_box_iou_quadri.py 3.21 KB
Newer Older
Yue Zhou's avatar
[Feature] Add box_iou_quadri & nms_quadri (#2277)  
Yue Zhou committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77 
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import pytest
import torch

from mmcv.utils import IS_CUDA_AVAILABLE


class TestBoxIoUQuadri:

    @pytest.mark.parametrize('device', [
        'cpu',
        pytest.param(
            'cuda',
            marks=pytest.mark.skipif(
                not IS_CUDA_AVAILABLE, reason='requires CUDA support')),
    ])
    def test_box_iou_quadri_cuda(self, device):
        from mmcv.ops import box_iou_quadri
        np_boxes1 = np.asarray([[1.0, 1.0, 3.0, 4.0, 4.0, 4.0, 4.0, 1.0],
                                [2.0, 2.0, 3.0, 4.0, 4.0, 2.0, 3.0, 1.0],
                                [7.0, 7.0, 8.0, 8.0, 9.0, 7.0, 8.0, 6.0]],
                               dtype=np.float32)
        np_boxes2 = np.asarray([[0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 0.0],
                                [2.0, 1.0, 2.0, 4.0, 4.0, 4.0, 4.0, 1.0],
                                [7.0, 6.0, 7.0, 8.0, 9.0, 8.0, 9.0, 6.0]],
                               dtype=np.float32)
        np_expect_ious = np.asarray(
            [[0.0714, 1.0000, 0.0000], [0.0000, 0.5000, 0.0000],
             [0.0000, 0.0000, 0.5000]],
            dtype=np.float32)
        np_expect_ious_aligned = np.asarray([0.0714, 0.5000, 0.5000],
                                            dtype=np.float32)

        boxes1 = torch.from_numpy(np_boxes1).to(device)
        boxes2 = torch.from_numpy(np_boxes2).to(device)

        ious = box_iou_quadri(boxes1, boxes2)
        assert np.allclose(ious.cpu().numpy(), np_expect_ious, atol=1e-4)

        ious = box_iou_quadri(boxes1, boxes2, aligned=True)
        assert np.allclose(
            ious.cpu().numpy(), np_expect_ious_aligned, atol=1e-4)

    @pytest.mark.parametrize('device', [
        'cpu',
        pytest.param(
            'cuda',
            marks=pytest.mark.skipif(
                not IS_CUDA_AVAILABLE, reason='requires CUDA support')),
    ])
    def test_box_iou_quadri_iof_cuda(self, device):
        from mmcv.ops import box_iou_quadri
        np_boxes1 = np.asarray([[1.0, 1.0, 3.0, 4.0, 4.0, 4.0, 4.0, 1.0],
                                [2.0, 2.0, 3.0, 4.0, 4.0, 2.0, 3.0, 1.0],
                                [7.0, 7.0, 8.0, 8.0, 9.0, 7.0, 8.0, 6.0]],
                               dtype=np.float32)
        np_boxes2 = np.asarray([[0.0, 0.0, 0.0, 2.0, 2.0, 2.0, 2.0, 0.0],
                                [2.0, 1.0, 2.0, 4.0, 4.0, 4.0, 4.0, 1.0],
                                [7.0, 6.0, 7.0, 8.0, 9.0, 8.0, 9.0, 6.0]],
                               dtype=np.float32)
        np_expect_ious = np.asarray(
            [[0.1111, 1.0000, 0.0000], [0.0000, 1.0000, 0.0000],
             [0.0000, 0.0000, 1.0000]],
            dtype=np.float32)
        np_expect_ious_aligned = np.asarray([0.1111, 1.0000, 1.0000],
                                            dtype=np.float32)

        boxes1 = torch.from_numpy(np_boxes1).to(device)
        boxes2 = torch.from_numpy(np_boxes2).to(device)

        ious = box_iou_quadri(boxes1, boxes2, mode='iof')
        assert np.allclose(ious.cpu().numpy(), np_expect_ious, atol=1e-4)

        ious = box_iou_quadri(boxes1, boxes2, mode='iof', aligned=True)
        assert np.allclose(
            ious.cpu().numpy(), np_expect_ious_aligned, atol=1e-4)