Spaces:
Sleeping
Sleeping
File size: 2,207 Bytes
d7e58f0 |
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 |
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import pytest
from numpy.testing import assert_array_equal
import mmcv
try:
import torch
except ImportError:
torch = None
@pytest.mark.skipif(torch is None, reason='requires torch library')
def test_tensor2imgs():
# test tensor obj
with pytest.raises(AssertionError):
tensor = np.random.rand(2, 3, 3)
mmcv.tensor2imgs(tensor)
# test tensor ndim
with pytest.raises(AssertionError):
tensor = torch.randn(2, 3, 3)
mmcv.tensor2imgs(tensor)
# test tensor dim-1
with pytest.raises(AssertionError):
tensor = torch.randn(2, 4, 3, 3)
mmcv.tensor2imgs(tensor)
# test mean length
with pytest.raises(AssertionError):
tensor = torch.randn(2, 3, 5, 5)
mmcv.tensor2imgs(tensor, mean=(1, ))
tensor = torch.randn(2, 1, 5, 5)
mmcv.tensor2imgs(tensor, mean=(0, 0, 0))
# test std length
with pytest.raises(AssertionError):
tensor = torch.randn(2, 3, 5, 5)
mmcv.tensor2imgs(tensor, std=(1, ))
tensor = torch.randn(2, 1, 5, 5)
mmcv.tensor2imgs(tensor, std=(1, 1, 1))
# test to_rgb
with pytest.raises(AssertionError):
tensor = torch.randn(2, 1, 5, 5)
mmcv.tensor2imgs(tensor, mean=(0, ), std=(1, ), to_rgb=True)
# test rgb=True
tensor = torch.randn(2, 3, 5, 5)
gts = [
t.cpu().numpy().transpose(1, 2, 0).astype(np.uint8)
for t in tensor.flip(1)
]
outputs = mmcv.tensor2imgs(tensor, to_rgb=True)
for gt, output in zip(gts, outputs):
assert_array_equal(gt, output)
# test rgb=False
tensor = torch.randn(2, 3, 5, 5)
gts = [t.cpu().numpy().transpose(1, 2, 0).astype(np.uint8) for t in tensor]
outputs = mmcv.tensor2imgs(tensor, to_rgb=False)
for gt, output in zip(gts, outputs):
assert_array_equal(gt, output)
# test tensor channel 1 and rgb=False
tensor = torch.randn(2, 1, 5, 5)
gts = [t.squeeze(0).cpu().numpy().astype(np.uint8) for t in tensor]
outputs = mmcv.tensor2imgs(tensor, to_rgb=False)
for gt, output in zip(gts, outputs):
assert_array_equal(gt, output)
|