[fbsync] Allow 2D numpy arrays as inputs for to_image (#8256)

Summary: Co-authored-by: Nicolas Hug <[email protected]> Reviewed By: vmoens Differential Revision: D55062777 fbshipit-source-id: 09475eeb59e2a31157e47e38ed978a4b0e86e29f
pytorch · Mar 20, 2024 · 18243ef · 18243ef
1 parent aa4968c
commit 18243ef
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/test/test_transforms_v2.py b/test/test_transforms_v2.py
@@ -5182,6 +5182,11 @@ def test_functional_and_transform(self, make_input, fn):
         if isinstance(input, torch.Tensor):
             assert output.data_ptr() == input.data_ptr()
 
+    def test_2d_np_array(self):
+        # Non-regression test for https://github.com/pytorch/vision/issues/8255
+        input = np.random.rand(10, 10)
+        assert F.to_image(input).shape == (1, 10, 10)
+
     def test_functional_error(self):
         with pytest.raises(TypeError, match="Input can either be a pure Tensor, a numpy array, or a PIL image"):
             F.to_image(object())

diff --git a/torchvision/transforms/v2/functional/_type_conversion.py b/torchvision/transforms/v2/functional/_type_conversion.py
@@ -11,7 +11,7 @@
 def to_image(inpt: Union[torch.Tensor, PIL.Image.Image, np.ndarray]) -> tv_tensors.Image:
     """See :class:`~torchvision.transforms.v2.ToImage` for details."""
     if isinstance(inpt, np.ndarray):
-        output = torch.from_numpy(inpt).permute((2, 0, 1)).contiguous()
+        output = torch.from_numpy(np.atleast_3d(inpt)).permute((2, 0, 1)).contiguous()
     elif isinstance(inpt, PIL.Image.Image):
         output = pil_to_tensor(inpt)
     elif isinstance(inpt, torch.Tensor):