|
- # Copyright (c) OpenMMLab. All rights reserved.
- import numpy as np
- import pytest
- import torch
-
- from mmdet.core.bbox import distance2bbox
- from mmdet.core.mask.structures import BitmapMasks, PolygonMasks
- from mmdet.core.utils import (center_of_mass, filter_scores_and_topk,
- flip_tensor, mask2ndarray, select_single_mlvl)
-
-
- def dummy_raw_polygon_masks(size):
- """
- Args:
- size (tuple): expected shape of dummy masks, (N, H, W)
-
- Return:
- list[list[ndarray]]: dummy mask
- """
- num_obj, height, width = size
- polygons = []
- for _ in range(num_obj):
- num_points = np.random.randint(5) * 2 + 6
- polygons.append([np.random.uniform(0, min(height, width), num_points)])
- return polygons
-
-
- def test_mask2ndarray():
- raw_masks = np.ones((3, 28, 28))
- bitmap_mask = BitmapMasks(raw_masks, 28, 28)
- output_mask = mask2ndarray(bitmap_mask)
- assert np.allclose(raw_masks, output_mask)
-
- raw_masks = dummy_raw_polygon_masks((3, 28, 28))
- polygon_masks = PolygonMasks(raw_masks, 28, 28)
- output_mask = mask2ndarray(polygon_masks)
- assert output_mask.shape == (3, 28, 28)
-
- raw_masks = np.ones((3, 28, 28))
- output_mask = mask2ndarray(raw_masks)
- assert np.allclose(raw_masks, output_mask)
-
- raw_masks = torch.ones((3, 28, 28))
- output_mask = mask2ndarray(raw_masks)
- assert np.allclose(raw_masks, output_mask)
-
- # test unsupported type
- raw_masks = []
- with pytest.raises(TypeError):
- output_mask = mask2ndarray(raw_masks)
-
-
- def test_distance2bbox():
- point = torch.Tensor([[74., 61.], [-29., 106.], [138., 61.], [29., 170.]])
-
- distance = torch.Tensor([[0., 0, 1., 1.], [1., 2., 10., 6.],
- [22., -29., 138., 61.], [54., -29., 170., 61.]])
- expected_decode_bboxes = torch.Tensor([[74., 61., 75., 62.],
- [0., 104., 0., 112.],
- [100., 90., 100., 120.],
- [0., 120., 100., 120.]])
- out_bbox = distance2bbox(point, distance, max_shape=(120, 100))
- assert expected_decode_bboxes.allclose(out_bbox)
- out = distance2bbox(point, distance, max_shape=torch.Tensor((120, 100)))
- assert expected_decode_bboxes.allclose(out)
-
- batch_point = point.unsqueeze(0).repeat(2, 1, 1)
- batch_distance = distance.unsqueeze(0).repeat(2, 1, 1)
- batch_out = distance2bbox(
- batch_point, batch_distance, max_shape=(120, 100))[0]
- assert out.allclose(batch_out)
- batch_out = distance2bbox(
- batch_point, batch_distance, max_shape=[(120, 100), (120, 100)])[0]
- assert out.allclose(batch_out)
-
- batch_out = distance2bbox(point, batch_distance, max_shape=(120, 100))[0]
- assert out.allclose(batch_out)
-
- # test max_shape is not equal to batch
- with pytest.raises(AssertionError):
- distance2bbox(
- batch_point,
- batch_distance,
- max_shape=[(120, 100), (120, 100), (32, 32)])
-
- rois = torch.zeros((0, 4))
- deltas = torch.zeros((0, 4))
- out = distance2bbox(rois, deltas, max_shape=(120, 100))
- assert rois.shape == out.shape
-
- rois = torch.zeros((2, 0, 4))
- deltas = torch.zeros((2, 0, 4))
- out = distance2bbox(rois, deltas, max_shape=(120, 100))
- assert rois.shape == out.shape
-
-
- @pytest.mark.parametrize('mask', [
- torch.ones((28, 28)),
- torch.zeros((28, 28)),
- torch.rand(28, 28) > 0.5,
- torch.tensor([[0, 0, 0, 0], [0, 1, 1, 0], [0, 1, 1, 0], [0, 0, 0, 0]])
- ])
- def test_center_of_mass(mask):
- center_h, center_w = center_of_mass(mask)
- if mask.shape[0] == 4:
- assert center_h == 1.5
- assert center_w == 1.5
- assert isinstance(center_h, torch.Tensor) \
- and isinstance(center_w, torch.Tensor)
- assert 0 <= center_h <= 28 \
- and 0 <= center_w <= 28
-
-
- def test_flip_tensor():
- img = np.random.random((1, 3, 10, 10))
- src_tensor = torch.from_numpy(img)
-
- # test flip_direction parameter error
- with pytest.raises(AssertionError):
- flip_tensor(src_tensor, 'flip')
-
- # test tensor dimension
- with pytest.raises(AssertionError):
- flip_tensor(src_tensor[0], 'vertical')
-
- hfilp_tensor = flip_tensor(src_tensor, 'horizontal')
- expected_hflip_tensor = torch.from_numpy(img[..., ::-1, :].copy())
- expected_hflip_tensor.allclose(hfilp_tensor)
-
- vfilp_tensor = flip_tensor(src_tensor, 'vertical')
- expected_vflip_tensor = torch.from_numpy(img[..., ::-1].copy())
- expected_vflip_tensor.allclose(vfilp_tensor)
-
- diag_filp_tensor = flip_tensor(src_tensor, 'diagonal')
- expected_diag_filp_tensor = torch.from_numpy(img[..., ::-1, ::-1].copy())
- expected_diag_filp_tensor.allclose(diag_filp_tensor)
-
-
- def test_select_single_mlvl():
- mlvl_tensors = [torch.rand(2, 1, 10, 10)] * 5
- mlvl_tensor_list = select_single_mlvl(mlvl_tensors, 1)
- assert len(mlvl_tensor_list) == 5 and mlvl_tensor_list[0].ndim == 3
-
-
- def test_filter_scores_and_topk():
- score = torch.tensor([[0.1, 0.3, 0.2], [0.12, 0.7, 0.9], [0.02, 0.8, 0.08],
- [0.4, 0.1, 0.08]])
- bbox_pred = torch.tensor([[0.2, 0.3], [0.4, 0.7], [0.1, 0.1], [0.5, 0.1]])
- score_thr = 0.15
- nms_pre = 4
- # test results type error
- with pytest.raises(NotImplementedError):
- filter_scores_and_topk(score, score_thr, nms_pre, (score, ))
-
- filtered_results = filter_scores_and_topk(
- score, score_thr, nms_pre, results=dict(bbox_pred=bbox_pred))
- filtered_score, labels, keep_idxs, results = filtered_results
- assert filtered_score.allclose(torch.tensor([0.9, 0.8, 0.7, 0.4]))
- assert labels.allclose(torch.tensor([2, 1, 1, 0]))
- assert keep_idxs.allclose(torch.tensor([1, 2, 1, 3]))
- assert results['bbox_pred'].allclose(
- torch.tensor([[0.4, 0.7], [0.1, 0.1], [0.4, 0.7], [0.5, 0.1]]))
|