File size: 1,253 Bytes
2366e36
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# Copyright (c) OpenMMLab. All rights reserved.
import numpy as np
import pytest
import torch
from numpy.testing import assert_array_equal

from mmocr.apis.utils import tensor2grayimgs


@pytest.mark.skipif(not torch.cuda.is_available(), reason='requires cuda')
def test_tensor2grayimgs():

    # test tensor obj
    with pytest.raises(AssertionError):
        tensor = np.random.rand(2, 3, 3)
        tensor2grayimgs(tensor)

    # test tensor ndim
    with pytest.raises(AssertionError):
        tensor = torch.randn(2, 3, 3)
        tensor2grayimgs(tensor)

    # test tensor dim-1
    with pytest.raises(AssertionError):
        tensor = torch.randn(2, 3, 5, 5)
        tensor2grayimgs(tensor)

    # test mean length
    with pytest.raises(AssertionError):
        tensor = torch.randn(2, 1, 5, 5)
        tensor2grayimgs(tensor, mean=(1, 1, 1))

    # test std length
    with pytest.raises(AssertionError):
        tensor = torch.randn(2, 1, 5, 5)
        tensor2grayimgs(tensor, std=(1, 1, 1))

    tensor = torch.randn(2, 1, 5, 5)
    gts = [t.squeeze(0).cpu().numpy().astype(np.uint8) for t in tensor]
    outputs = tensor2grayimgs(tensor, mean=(0, ), std=(1, ))
    for gt, output in zip(gts, outputs):
        assert_array_equal(gt, output)