| 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): |
| |
| 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)) |
|
|
