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

Summary: Co-authored-by: Nicolas Hug <contact@nicolas-hug.com> Reviewed By: vmoens Differential Revision: D55062777 fbshipit-source-id: 09475eeb59e2a31157e47e38ed978a4b0e86e29f

Author: NicolasHug

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