File size: 1,314 Bytes
c508d7f |
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 |
from tensorboardX import summary
from tensorboardX.utils import make_grid, _prepare_video, convert_to_HWC
import numpy as np
import pytest
import unittest
class UtilsTest(unittest.TestCase):
def test_to_HWC(self):
np.random.seed(1)
test_image = np.random.randint(0, 256, size=(3, 32, 32), dtype=np.uint8)
converted = convert_to_HWC(test_image, 'chw')
assert converted.shape == (32, 32, 3)
test_image = np.random.randint(0, 256, size=(16, 3, 32, 32), dtype=np.uint8)
converted = convert_to_HWC(test_image, 'nchw')
assert converted.shape == (64, 256, 3)
test_image = np.random.randint(0, 256, size=(32, 32), dtype=np.uint8)
converted = convert_to_HWC(test_image, 'hw')
assert converted.shape == (32, 32, 3)
def test_prepare_video(self):
# at each timestep the sum over all other dimensions of the video should stay the same
np.random.seed(1)
V_before = np.random.random((4, 10, 3, 20, 20))
V_after = _prepare_video(np.copy(V_before))
V_before = np.swapaxes(V_before, 0, 1)
V_before = np.reshape(V_before, newshape=(10, -1))
V_after = np.reshape(V_after, newshape=(10, -1))
np.testing.assert_array_almost_equal(np.sum(V_before, axis=1), np.sum(V_after, axis=1))
|