diff --git a/test/test_transforms.py b/test/test_transforms.py index 8a0fb0f5ca6..e86b6959517 100644 --- a/test/test_transforms.py +++ b/test/test_transforms.py @@ -195,100 +195,98 @@ def test_accimage_crop(self): torch.testing.assert_close(output, expected_output) -@pytest.mark.parametrize('channels', [1, 3, 4]) -def test_to_tensor(channels): - height, width = 4, 4 - trans = transforms.ToTensor() - - input_data = torch.ByteTensor(channels, height, width).random_(0, 255).float().div_(255) - img = transforms.ToPILImage()(input_data) - output = trans(img) - torch.testing.assert_close(output, input_data, check_stride=False) - - ndarray = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8) - output = trans(ndarray) - expected_output = ndarray.transpose((2, 0, 1)) / 255.0 - torch.testing.assert_close(output.numpy(), expected_output, check_stride=False, check_dtype=False) - - ndarray = np.random.rand(height, width, channels).astype(np.float32) - output = trans(ndarray) - expected_output = ndarray.transpose((2, 0, 1)) - torch.testing.assert_close(output.numpy(), expected_output, check_stride=False, check_dtype=False) - - # separate test for mode '1' PIL images - input_data = torch.ByteTensor(1, height, width).bernoulli_() - img = transforms.ToPILImage()(input_data.mul(255)).convert('1') - output = trans(img) - torch.testing.assert_close(input_data, output, check_dtype=False, check_stride=False) - - -def test_to_tensor_errors(): - height, width = 4, 4 - trans = transforms.ToTensor() - - with pytest.raises(TypeError): - trans(np.random.rand(1, height, width).tolist()) - - with pytest.raises(ValueError): - trans(np.random.rand(height)) - - with pytest.raises(ValueError): - trans(np.random.rand(1, 1, height, width)) +class TestToTensor: + @pytest.mark.parametrize('channels', [1, 3, 4]) + def test_to_tensor(self, channels): + height, width = 4, 4 + trans = transforms.ToTensor() -@pytest.mark.parametrize('dtype', [torch.float16, torch.float, torch.double]) -def test_to_tensor_with_other_default_dtypes(dtype): - current_def_dtype = torch.get_default_dtype() + input_data = torch.ByteTensor(channels, height, width).random_(0, 255).float().div_(255) + img = transforms.ToPILImage()(input_data) + output = trans(img) + torch.testing.assert_close(output, input_data, check_stride=False) + + ndarray = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8) + output = trans(ndarray) + expected_output = ndarray.transpose((2, 0, 1)) / 255.0 + torch.testing.assert_close(output.numpy(), expected_output, check_stride=False, check_dtype=False) + + ndarray = np.random.rand(height, width, channels).astype(np.float32) + output = trans(ndarray) + expected_output = ndarray.transpose((2, 0, 1)) + torch.testing.assert_close(output.numpy(), expected_output, check_stride=False, check_dtype=False) + + # separate test for mode '1' PIL images + input_data = torch.ByteTensor(1, height, width).bernoulli_() + img = transforms.ToPILImage()(input_data.mul(255)).convert('1') + output = trans(img) + torch.testing.assert_close(input_data, output, check_dtype=False, check_stride=False) + + def test_to_tensor_errors(self): + height, width = 4, 4 + trans = transforms.ToTensor() - t = transforms.ToTensor() - np_arr = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8) - img = Image.fromarray(np_arr) + with pytest.raises(TypeError): + trans(np.random.rand(1, height, width).tolist()) - torch.set_default_dtype(dtype) - res = t(img) - assert res.dtype == dtype, f"{res.dtype} vs {dtype}" + with pytest.raises(ValueError): + trans(np.random.rand(height)) - torch.set_default_dtype(current_def_dtype) + with pytest.raises(ValueError): + trans(np.random.rand(1, 1, height, width)) + @pytest.mark.parametrize('dtype', [torch.float16, torch.float, torch.double]) + def test_to_tensor_with_other_default_dtypes(self, dtype): + current_def_dtype = torch.get_default_dtype() -@pytest.mark.parametrize('channels', [1, 3, 4]) -def test_pil_to_tensor(channels): - height, width = 4, 4 - trans = transforms.PILToTensor() + t = transforms.ToTensor() + np_arr = np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8) + img = Image.fromarray(np_arr) - input_data = torch.ByteTensor(channels, height, width).random_(0, 255) - img = transforms.ToPILImage()(input_data) - output = trans(img) - torch.testing.assert_close(input_data, output, check_stride=False) + torch.set_default_dtype(dtype) + res = t(img) + assert res.dtype == dtype, f"{res.dtype} vs {dtype}" - input_data = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8) - img = transforms.ToPILImage()(input_data) - output = trans(img) - expected_output = input_data.transpose((2, 0, 1)) - torch.testing.assert_close(output.numpy(), expected_output) + torch.set_default_dtype(current_def_dtype) - input_data = torch.as_tensor(np.random.rand(channels, height, width).astype(np.float32)) - img = transforms.ToPILImage()(input_data) # CHW -> HWC and (* 255).byte() - output = trans(img) # HWC -> CHW - expected_output = (input_data * 255).byte() - torch.testing.assert_close(output, expected_output, check_stride=False) + @pytest.mark.parametrize('channels', [1, 3, 4]) + def test_pil_to_tensor(self, channels): + height, width = 4, 4 + trans = transforms.PILToTensor() - # separate test for mode '1' PIL images - input_data = torch.ByteTensor(1, height, width).bernoulli_() - img = transforms.ToPILImage()(input_data.mul(255)).convert('1') - output = trans(img).view(torch.uint8).bool().to(torch.uint8) - torch.testing.assert_close(input_data, output, check_stride=False) + input_data = torch.ByteTensor(channels, height, width).random_(0, 255) + img = transforms.ToPILImage()(input_data) + output = trans(img) + torch.testing.assert_close(input_data, output, check_stride=False) + + input_data = np.random.randint(low=0, high=255, size=(height, width, channels)).astype(np.uint8) + img = transforms.ToPILImage()(input_data) + output = trans(img) + expected_output = input_data.transpose((2, 0, 1)) + torch.testing.assert_close(output.numpy(), expected_output) + + input_data = torch.as_tensor(np.random.rand(channels, height, width).astype(np.float32)) + img = transforms.ToPILImage()(input_data) # CHW -> HWC and (* 255).byte() + output = trans(img) # HWC -> CHW + expected_output = (input_data * 255).byte() + torch.testing.assert_close(output, expected_output, check_stride=False) + # separate test for mode '1' PIL images + input_data = torch.ByteTensor(1, height, width).bernoulli_() + img = transforms.ToPILImage()(input_data.mul(255)).convert('1') + output = trans(img).view(torch.uint8).bool().to(torch.uint8) + torch.testing.assert_close(input_data, output, check_stride=False) -def test_pil_to_tensor_errors(): - height, width = 4, 4 - trans = transforms.PILToTensor() + def test_pil_to_tensor_errors(self): + height, width = 4, 4 + trans = transforms.PILToTensor() - with pytest.raises(TypeError): - trans(np.random.rand(1, height, width).tolist()) + with pytest.raises(TypeError): + trans(np.random.rand(1, height, width).tolist()) - with pytest.raises(TypeError): - trans(np.random.rand(1, height, width)) + with pytest.raises(TypeError): + trans(np.random.rand(1, height, width)) def test_randomresized_params(): @@ -537,303 +535,285 @@ def test_randomness(fn, trans, config, p): assert p_value > 0.0001 -def _get_1_channel_tensor_various_types(): - img_data_float = torch.Tensor(1, 4, 4).uniform_() - expected_output = img_data_float.mul(255).int().float().div(255).numpy() - yield img_data_float, expected_output, 'L' +class TestToPil: - img_data_byte = torch.ByteTensor(1, 4, 4).random_(0, 255) - expected_output = img_data_byte.float().div(255.0).numpy() - yield img_data_byte, expected_output, 'L' + def _get_1_channel_tensor_various_types(): + img_data_float = torch.Tensor(1, 4, 4).uniform_() + expected_output = img_data_float.mul(255).int().float().div(255).numpy() + yield img_data_float, expected_output, 'L' - img_data_short = torch.ShortTensor(1, 4, 4).random_() - expected_output = img_data_short.numpy() - yield img_data_short, expected_output, 'I;16' + img_data_byte = torch.ByteTensor(1, 4, 4).random_(0, 255) + expected_output = img_data_byte.float().div(255.0).numpy() + yield img_data_byte, expected_output, 'L' - img_data_int = torch.IntTensor(1, 4, 4).random_() - expected_output = img_data_int.numpy() - yield img_data_int, expected_output, 'I' + img_data_short = torch.ShortTensor(1, 4, 4).random_() + expected_output = img_data_short.numpy() + yield img_data_short, expected_output, 'I;16' + img_data_int = torch.IntTensor(1, 4, 4).random_() + expected_output = img_data_int.numpy() + yield img_data_int, expected_output, 'I' -@pytest.mark.parametrize('with_mode', [False, True]) -@pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_1_channel_tensor_various_types()) -def test_1_channel_tensor_to_pil_image(with_mode, img_data, expected_output, expected_mode): - transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() - to_tensor = transforms.ToTensor() + def _get_2d_tensor_various_types(): + img_data_float = torch.Tensor(4, 4).uniform_() + expected_output = img_data_float.mul(255).int().float().div(255).numpy() + yield img_data_float, expected_output, 'L' - img = transform(img_data) - assert img.mode == expected_mode - torch.testing.assert_close(expected_output, to_tensor(img).numpy(), check_stride=False) + img_data_byte = torch.ByteTensor(4, 4).random_(0, 255) + expected_output = img_data_byte.float().div(255.0).numpy() + yield img_data_byte, expected_output, 'L' + img_data_short = torch.ShortTensor(4, 4).random_() + expected_output = img_data_short.numpy() + yield img_data_short, expected_output, 'I;16' -def test_1_channel_float_tensor_to_pil_image(): - img_data = torch.Tensor(1, 4, 4).uniform_() - # 'F' mode for torch.FloatTensor - img_F_mode = transforms.ToPILImage(mode='F')(img_data) - assert img_F_mode.mode == 'F' - torch.testing.assert_close( - np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode='F')), np.array(img_F_mode) - ) + img_data_int = torch.IntTensor(4, 4).random_() + expected_output = img_data_int.numpy() + yield img_data_int, expected_output, 'I' + @pytest.mark.parametrize('with_mode', [False, True]) + @pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_1_channel_tensor_various_types()) + def test_1_channel_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode): + transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() + to_tensor = transforms.ToTensor() -@pytest.mark.parametrize('with_mode', [False, True]) -@pytest.mark.parametrize('img_data, expected_mode', [ - (torch.Tensor(4, 4, 1).uniform_().numpy(), 'F'), - (torch.ByteTensor(4, 4, 1).random_(0, 255).numpy(), 'L'), - (torch.ShortTensor(4, 4, 1).random_().numpy(), 'I;16'), - (torch.IntTensor(4, 4, 1).random_().numpy(), 'I'), -]) -def test_1_channel_ndarray_to_pil_image(with_mode, img_data, expected_mode): - transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() - img = transform(img_data) - assert img.mode == expected_mode - # note: we explicitly convert img's dtype because pytorch doesn't support uint16 - # and otherwise assert_close wouldn't be able to construct a tensor from the uint16 array - torch.testing.assert_close(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype)) - - -@pytest.mark.parametrize('expected_mode', [None, 'LA']) -def test_2_channel_ndarray_to_pil_image(expected_mode): - img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy() - - if expected_mode is None: - img = transforms.ToPILImage()(img_data) - assert img.mode == 'LA' # default should assume LA - else: - img = transforms.ToPILImage(mode=expected_mode)(img_data) - assert img.mode == expected_mode - split = img.split() - for i in range(2): - torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False) - - -def test_2_channel_ndarray_to_pil_image_error(): - img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy() - transforms.ToPILImage().__repr__() - - # should raise if we try a mode for 4 or 1 or 3 channel images - with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): - transforms.ToPILImage(mode='RGBA')(img_data) - with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): - transforms.ToPILImage(mode='P')(img_data) - with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): - transforms.ToPILImage(mode='RGB')(img_data) - - -@pytest.mark.parametrize('expected_mode', [None, 'LA']) -def test_2_channel_tensor_to_pil_image(expected_mode): - img_data = torch.Tensor(2, 4, 4).uniform_() - expected_output = img_data.mul(255).int().float().div(255) - if expected_mode is None: - img = transforms.ToPILImage()(img_data) - assert img.mode == 'LA' # default should assume LA - else: - img = transforms.ToPILImage(mode=expected_mode)(img_data) + img = transform(img_data) assert img.mode == expected_mode + torch.testing.assert_close(expected_output, to_tensor(img).numpy(), check_stride=False) - split = img.split() - for i in range(2): - torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) - - -def test_2_channel_tensor_to_pil_image_error(): - img_data = torch.Tensor(2, 4, 4).uniform_() - - # should raise if we try a mode for 4 or 1 or 3 channel images - with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): - transforms.ToPILImage(mode='RGBA')(img_data) - with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): - transforms.ToPILImage(mode='P')(img_data) - with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): - transforms.ToPILImage(mode='RGB')(img_data) - - -def _get_2d_tensor_various_types(): - img_data_float = torch.Tensor(4, 4).uniform_() - expected_output = img_data_float.mul(255).int().float().div(255).numpy() - yield img_data_float, expected_output, 'L' - - img_data_byte = torch.ByteTensor(4, 4).random_(0, 255) - expected_output = img_data_byte.float().div(255.0).numpy() - yield img_data_byte, expected_output, 'L' - - img_data_short = torch.ShortTensor(4, 4).random_() - expected_output = img_data_short.numpy() - yield img_data_short, expected_output, 'I;16' - - img_data_int = torch.IntTensor(4, 4).random_() - expected_output = img_data_int.numpy() - yield img_data_int, expected_output, 'I' - - -@pytest.mark.parametrize('with_mode', [False, True]) -@pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_2d_tensor_various_types()) -def test_2d_tensor_to_pil_image(with_mode, img_data, expected_output, expected_mode): - transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() - to_tensor = transforms.ToTensor() - - img = transform(img_data) - assert img.mode == expected_mode - torch.testing.assert_close(expected_output, to_tensor(img).numpy()[0]) - - -@pytest.mark.parametrize('with_mode', [False, True]) -@pytest.mark.parametrize('img_data, expected_mode', [ - (torch.Tensor(4, 4).uniform_().numpy(), 'F'), - (torch.ByteTensor(4, 4).random_(0, 255).numpy(), 'L'), - (torch.ShortTensor(4, 4).random_().numpy(), 'I;16'), - (torch.IntTensor(4, 4).random_().numpy(), 'I'), -]) -def test_2d_ndarray_to_pil_image(with_mode, img_data, expected_mode): - transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() - img = transform(img_data) - assert img.mode == expected_mode - np.testing.assert_allclose(img_data, img) - - -@pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr']) -def test_3_channel_tensor_to_pil_image(expected_mode): - img_data = torch.Tensor(3, 4, 4).uniform_() - expected_output = img_data.mul(255).int().float().div(255) + def test_1_channel_float_tensor_to_pil_image(self): + img_data = torch.Tensor(1, 4, 4).uniform_() + # 'F' mode for torch.FloatTensor + img_F_mode = transforms.ToPILImage(mode='F')(img_data) + assert img_F_mode.mode == 'F' + torch.testing.assert_close( + np.array(Image.fromarray(img_data.squeeze(0).numpy(), mode='F')), np.array(img_F_mode) + ) - if expected_mode is None: - img = transforms.ToPILImage()(img_data) - assert img.mode == 'RGB' # default should assume RGB - else: - img = transforms.ToPILImage(mode=expected_mode)(img_data) + @pytest.mark.parametrize('with_mode', [False, True]) + @pytest.mark.parametrize('img_data, expected_mode', [ + (torch.Tensor(4, 4, 1).uniform_().numpy(), 'F'), + (torch.ByteTensor(4, 4, 1).random_(0, 255).numpy(), 'L'), + (torch.ShortTensor(4, 4, 1).random_().numpy(), 'I;16'), + (torch.IntTensor(4, 4, 1).random_().numpy(), 'I'), + ]) + def test_1_channel_ndarray_to_pil_image(self, with_mode, img_data, expected_mode): + transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() + img = transform(img_data) assert img.mode == expected_mode - split = img.split() - for i in range(3): - torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) + # note: we explicitly convert img's dtype because pytorch doesn't support uint16 + # and otherwise assert_close wouldn't be able to construct a tensor from the uint16 array + torch.testing.assert_close(img_data[:, :, 0], np.asarray(img).astype(img_data.dtype)) + @pytest.mark.parametrize('expected_mode', [None, 'LA']) + def test_2_channel_ndarray_to_pil_image(self, expected_mode): + img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy() -def test_3_channel_tensor_to_pil_image_error(): - img_data = torch.Tensor(3, 4, 4).uniform_() - error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs" - # should raise if we try a mode for 4 or 1 or 2 channel images - with pytest.raises(ValueError, match=error_message_3d): - transforms.ToPILImage(mode='RGBA')(img_data) - with pytest.raises(ValueError, match=error_message_3d): - transforms.ToPILImage(mode='P')(img_data) - with pytest.raises(ValueError, match=error_message_3d): - transforms.ToPILImage(mode='LA')(img_data) - - with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'): - transforms.ToPILImage()(torch.Tensor(1, 3, 4, 4).uniform_()) - - -@pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr']) -def test_3_channel_ndarray_to_pil_image(expected_mode): - img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy() - - if expected_mode is None: - img = transforms.ToPILImage()(img_data) - assert img.mode == 'RGB' # default should assume RGB - else: - img = transforms.ToPILImage(mode=expected_mode)(img_data) + if expected_mode is None: + img = transforms.ToPILImage()(img_data) + assert img.mode == 'LA' # default should assume LA + else: + img = transforms.ToPILImage(mode=expected_mode)(img_data) + assert img.mode == expected_mode + split = img.split() + for i in range(2): + torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False) + + def test_2_channel_ndarray_to_pil_image_error(self): + img_data = torch.ByteTensor(4, 4, 2).random_(0, 255).numpy() + transforms.ToPILImage().__repr__() + + # should raise if we try a mode for 4 or 1 or 3 channel images + with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): + transforms.ToPILImage(mode='RGBA')(img_data) + with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): + transforms.ToPILImage(mode='P')(img_data) + with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): + transforms.ToPILImage(mode='RGB')(img_data) + + @pytest.mark.parametrize('expected_mode', [None, 'LA']) + def test_2_channel_tensor_to_pil_image(self, expected_mode): + img_data = torch.Tensor(2, 4, 4).uniform_() + expected_output = img_data.mul(255).int().float().div(255) + if expected_mode is None: + img = transforms.ToPILImage()(img_data) + assert img.mode == 'LA' # default should assume LA + else: + img = transforms.ToPILImage(mode=expected_mode)(img_data) + assert img.mode == expected_mode + + split = img.split() + for i in range(2): + torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) + + def test_2_channel_tensor_to_pil_image_error(self): + img_data = torch.Tensor(2, 4, 4).uniform_() + + # should raise if we try a mode for 4 or 1 or 3 channel images + with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): + transforms.ToPILImage(mode='RGBA')(img_data) + with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): + transforms.ToPILImage(mode='P')(img_data) + with pytest.raises(ValueError, match=r"Only modes \['LA'\] are supported for 2D inputs"): + transforms.ToPILImage(mode='RGB')(img_data) + + @pytest.mark.parametrize('with_mode', [False, True]) + @pytest.mark.parametrize('img_data, expected_output, expected_mode', _get_2d_tensor_various_types()) + def test_2d_tensor_to_pil_image(self, with_mode, img_data, expected_output, expected_mode): + transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() + to_tensor = transforms.ToTensor() + + img = transform(img_data) assert img.mode == expected_mode - split = img.split() - for i in range(3): - torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False) - - -def test_3_channel_ndarray_to_pil_image_error(): - img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy() - - # Checking if ToPILImage can be printed as string - transforms.ToPILImage().__repr__() - - error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs" - # should raise if we try a mode for 4 or 1 or 2 channel images - with pytest.raises(ValueError, match=error_message_3d): - transforms.ToPILImage(mode='RGBA')(img_data) - with pytest.raises(ValueError, match=error_message_3d): - transforms.ToPILImage(mode='P')(img_data) - with pytest.raises(ValueError, match=error_message_3d): - transforms.ToPILImage(mode='LA')(img_data) - - -@pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX']) -def test_4_channel_tensor_to_pil_image(expected_mode): - img_data = torch.Tensor(4, 4, 4).uniform_() - expected_output = img_data.mul(255).int().float().div(255) - - if expected_mode is None: - img = transforms.ToPILImage()(img_data) - assert img.mode == 'RGBA' # default should assume RGBA - else: - img = transforms.ToPILImage(mode=expected_mode)(img_data) + torch.testing.assert_close(expected_output, to_tensor(img).numpy()[0]) + + @pytest.mark.parametrize('with_mode', [False, True]) + @pytest.mark.parametrize('img_data, expected_mode', [ + (torch.Tensor(4, 4).uniform_().numpy(), 'F'), + (torch.ByteTensor(4, 4).random_(0, 255).numpy(), 'L'), + (torch.ShortTensor(4, 4).random_().numpy(), 'I;16'), + (torch.IntTensor(4, 4).random_().numpy(), 'I'), + ]) + def test_2d_ndarray_to_pil_image(self, with_mode, img_data, expected_mode): + transform = transforms.ToPILImage(mode=expected_mode) if with_mode else transforms.ToPILImage() + img = transform(img_data) assert img.mode == expected_mode + np.testing.assert_allclose(img_data, img) - split = img.split() - for i in range(4): - torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) - - -def test_4_channel_tensor_to_pil_image_error(): - img_data = torch.Tensor(4, 4, 4).uniform_() - - error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs" - # should raise if we try a mode for 3 or 1 or 2 channel images - with pytest.raises(ValueError, match=error_message_4d): - transforms.ToPILImage(mode='RGB')(img_data) - with pytest.raises(ValueError, match=error_message_4d): - transforms.ToPILImage(mode='P')(img_data) - with pytest.raises(ValueError, match=error_message_4d): - transforms.ToPILImage(mode='LA')(img_data) + @pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr']) + def test_3_channel_tensor_to_pil_image(self, expected_mode): + img_data = torch.Tensor(3, 4, 4).uniform_() + expected_output = img_data.mul(255).int().float().div(255) - -@pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX']) -def test_4_channel_ndarray_to_pil_image(expected_mode): - img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy() - - if expected_mode is None: - img = transforms.ToPILImage()(img_data) - assert img.mode == 'RGBA' # default should assume RGBA - else: - img = transforms.ToPILImage(mode=expected_mode)(img_data) - assert img.mode == expected_mode - split = img.split() - for i in range(4): - torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False) - - -def test_4_channel_ndarray_to_pil_image_error(): - img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy() - - error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs" - # should raise if we try a mode for 3 or 1 or 2 channel images - with pytest.raises(ValueError, match=error_message_4d): - transforms.ToPILImage(mode='RGB')(img_data) - with pytest.raises(ValueError, match=error_message_4d): - transforms.ToPILImage(mode='P')(img_data) - with pytest.raises(ValueError, match=error_message_4d): - transforms.ToPILImage(mode='LA')(img_data) - - -def test_ndarray_bad_types_to_pil_image(): - trans = transforms.ToPILImage() - reg_msg = r'Input type \w+ is not supported' - with pytest.raises(TypeError, match=reg_msg): - trans(np.ones([4, 4, 1], np.int64)) - with pytest.raises(TypeError, match=reg_msg): - trans(np.ones([4, 4, 1], np.uint16)) - with pytest.raises(TypeError, match=reg_msg): - trans(np.ones([4, 4, 1], np.uint32)) - with pytest.raises(TypeError, match=reg_msg): - trans(np.ones([4, 4, 1], np.float64)) - - with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'): - transforms.ToPILImage()(np.ones([1, 4, 4, 3])) - with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'): - transforms.ToPILImage()(np.ones([4, 4, 6])) - - -def test_tensor_bad_types_to_pil_image(): - with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'): - transforms.ToPILImage()(torch.ones(1, 3, 4, 4)) - with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'): - transforms.ToPILImage()(torch.ones(6, 4, 4)) + if expected_mode is None: + img = transforms.ToPILImage()(img_data) + assert img.mode == 'RGB' # default should assume RGB + else: + img = transforms.ToPILImage(mode=expected_mode)(img_data) + assert img.mode == expected_mode + split = img.split() + for i in range(3): + torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) + + def test_3_channel_tensor_to_pil_image_error(self): + img_data = torch.Tensor(3, 4, 4).uniform_() + error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs" + # should raise if we try a mode for 4 or 1 or 2 channel images + with pytest.raises(ValueError, match=error_message_3d): + transforms.ToPILImage(mode='RGBA')(img_data) + with pytest.raises(ValueError, match=error_message_3d): + transforms.ToPILImage(mode='P')(img_data) + with pytest.raises(ValueError, match=error_message_3d): + transforms.ToPILImage(mode='LA')(img_data) + + with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'): + transforms.ToPILImage()(torch.Tensor(1, 3, 4, 4).uniform_()) + + @pytest.mark.parametrize('expected_mode', [None, 'RGB', 'HSV', 'YCbCr']) + def test_3_channel_ndarray_to_pil_image(self, expected_mode): + img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy() + + if expected_mode is None: + img = transforms.ToPILImage()(img_data) + assert img.mode == 'RGB' # default should assume RGB + else: + img = transforms.ToPILImage(mode=expected_mode)(img_data) + assert img.mode == expected_mode + split = img.split() + for i in range(3): + torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False) + + def test_3_channel_ndarray_to_pil_image_error(self): + img_data = torch.ByteTensor(4, 4, 3).random_(0, 255).numpy() + + # Checking if ToPILImage can be printed as string + transforms.ToPILImage().__repr__() + + error_message_3d = r"Only modes \['RGB', 'YCbCr', 'HSV'\] are supported for 3D inputs" + # should raise if we try a mode for 4 or 1 or 2 channel images + with pytest.raises(ValueError, match=error_message_3d): + transforms.ToPILImage(mode='RGBA')(img_data) + with pytest.raises(ValueError, match=error_message_3d): + transforms.ToPILImage(mode='P')(img_data) + with pytest.raises(ValueError, match=error_message_3d): + transforms.ToPILImage(mode='LA')(img_data) + + @pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX']) + def test_4_channel_tensor_to_pil_image(self, expected_mode): + img_data = torch.Tensor(4, 4, 4).uniform_() + expected_output = img_data.mul(255).int().float().div(255) + + if expected_mode is None: + img = transforms.ToPILImage()(img_data) + assert img.mode == 'RGBA' # default should assume RGBA + else: + img = transforms.ToPILImage(mode=expected_mode)(img_data) + assert img.mode == expected_mode + + split = img.split() + for i in range(4): + torch.testing.assert_close(expected_output[i].numpy(), F.to_tensor(split[i]).squeeze(0).numpy()) + + def test_4_channel_tensor_to_pil_image_error(self): + img_data = torch.Tensor(4, 4, 4).uniform_() + + error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs" + # should raise if we try a mode for 3 or 1 or 2 channel images + with pytest.raises(ValueError, match=error_message_4d): + transforms.ToPILImage(mode='RGB')(img_data) + with pytest.raises(ValueError, match=error_message_4d): + transforms.ToPILImage(mode='P')(img_data) + with pytest.raises(ValueError, match=error_message_4d): + transforms.ToPILImage(mode='LA')(img_data) + + @pytest.mark.parametrize('expected_mode', [None, 'RGBA', 'CMYK', 'RGBX']) + def test_4_channel_ndarray_to_pil_image(self, expected_mode): + img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy() + + if expected_mode is None: + img = transforms.ToPILImage()(img_data) + assert img.mode == 'RGBA' # default should assume RGBA + else: + img = transforms.ToPILImage(mode=expected_mode)(img_data) + assert img.mode == expected_mode + split = img.split() + for i in range(4): + torch.testing.assert_close(img_data[:, :, i], np.asarray(split[i]), check_stride=False) + + def test_4_channel_ndarray_to_pil_image_error(self): + img_data = torch.ByteTensor(4, 4, 4).random_(0, 255).numpy() + + error_message_4d = r"Only modes \['RGBA', 'CMYK', 'RGBX'\] are supported for 4D inputs" + # should raise if we try a mode for 3 or 1 or 2 channel images + with pytest.raises(ValueError, match=error_message_4d): + transforms.ToPILImage(mode='RGB')(img_data) + with pytest.raises(ValueError, match=error_message_4d): + transforms.ToPILImage(mode='P')(img_data) + with pytest.raises(ValueError, match=error_message_4d): + transforms.ToPILImage(mode='LA')(img_data) + + def test_ndarray_bad_types_to_pil_image(self): + trans = transforms.ToPILImage() + reg_msg = r'Input type \w+ is not supported' + with pytest.raises(TypeError, match=reg_msg): + trans(np.ones([4, 4, 1], np.int64)) + with pytest.raises(TypeError, match=reg_msg): + trans(np.ones([4, 4, 1], np.uint16)) + with pytest.raises(TypeError, match=reg_msg): + trans(np.ones([4, 4, 1], np.uint32)) + with pytest.raises(TypeError, match=reg_msg): + trans(np.ones([4, 4, 1], np.float64)) + + with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'): + transforms.ToPILImage()(np.ones([1, 4, 4, 3])) + with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'): + transforms.ToPILImage()(np.ones([4, 4, 6])) + + def test_tensor_bad_types_to_pil_image(self): + with pytest.raises(ValueError, match=r'pic should be 2/3 dimensional. Got \d+ dimensions.'): + transforms.ToPILImage()(torch.ones(1, 3, 4, 4)) + with pytest.raises(ValueError, match=r'pic should not have > 4 channels. Got \d+ channels.'): + transforms.ToPILImage()(torch.ones(6, 4, 4)) def test_adjust_brightness():