mmcv/tests/test_ops/test_bbox.py
Zaida Zhou 0b4285d959
Pick commits from master (#2164)
* [Docs] Add swith_language.md in docs (#2160)

* [Fix] Fix onnx unit tests (#2155)

* [Docs] Limit extension versions (#2144)

* Support PrRoIPool operation

* Add MPS bbox overlap

* Add .pre-commit-config-zh-cn.yaml (#2135)

Co-authored-by: xcnick <xcnick0412@gmail.com>
Co-authored-by: Jingwei Zhang <zjw18@mails.tsinghua.edu.cn>
Co-authored-by: q.yao <yaoqian@sensetime.com>
2022-08-03 14:11:53 +08:00

67 lines
2.7 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, IS_MPS_AVAILABLE
class TestBBox:
def _test_bbox_overlaps(self, device='cpu', dtype=torch.float):
from mmcv.ops import bbox_overlaps
b1 = torch.tensor([[1.0, 1.0, 3.0, 4.0], [2.0, 2.0, 3.0, 4.0],
[7.0, 7.0, 8.0, 8.0]]).to(device).type(dtype)
b2 = torch.tensor([[0.0, 2.0, 2.0, 5.0], [2.0, 1.0, 3.0,
3.0]]).to(device).type(dtype)
should_output = np.array([[0.33333334, 0.5], [0.2, 0.5], [0.0, 0.0]])
out = bbox_overlaps(b1, b2, offset=1)
assert np.allclose(out.cpu().numpy(), should_output, 1e-2)
b1 = torch.tensor([[1.0, 1.0, 3.0, 4.0], [2.0, 2.0, 3.0,
4.0]]).to(device).type(dtype)
b2 = torch.tensor([[0.0, 2.0, 2.0, 5.0], [2.0, 1.0, 3.0,
3.0]]).to(device).type(dtype)
should_output = np.array([0.33333334, 0.5])
out = bbox_overlaps(b1, b2, aligned=True, offset=1)
assert np.allclose(out.cpu().numpy(), should_output, 1e-2)
b1 = torch.tensor([[0.0, 0.0, 3.0, 3.0]]).to(device).type(dtype)
b2 = torch.tensor([[4.0, 0.0, 5.0, 3.0], [3.0, 0.0, 4.0, 3.0],
[2.0, 0.0, 3.0, 3.0], [1.0, 0.0, 2.0,
3.0]]).to(device).type(dtype)
should_output = np.array([0, 0.2, 0.5, 0.5])
out = bbox_overlaps(b1, b2, offset=1)
assert np.allclose(out.cpu().numpy(), should_output, 1e-2)
@pytest.mark.parametrize('device', [
'cpu',
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')),
pytest.param(
'mps',
marks=pytest.mark.skipif(
not IS_MPS_AVAILABLE, reason='requires MPS support'))
])
def test_bbox_overlaps_float(self, device):
self._test_bbox_overlaps(device, dtype=torch.float)
@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_bbox_overlaps_half(self, device):
self._test_bbox_overlaps(device, dtype=torch.half)