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

Co-authored-by: Nicolas Hug <[email protected]>

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

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