Update torch_formatter.py (#6402)

* Update torch_formatter.py

Transpose images to (C, H, W) instead of (H, W, C).

* handle monochromat PIL image

* Update test_formatting.py

---------

Co-authored-by: Quentin Lhoest <[email protected]>

Author: varunneal

src/datasets/formatting/torch_formatter.py

@@ -71,6 +71,10 @@ def _tensorize(self, value):
 
             if isinstance(value, PIL.Image.Image):
                 value = np.asarray(value)
+                if value.ndim == 2:
+                    value = value[:, :, np.newaxis]
+
+                value = value.transpose((2, 0, 1))
         return torch.tensor(value, **{**default_dtype, **self.torch_tensor_kwargs})
 
     def _recursive_tensorize(self, data_struct):

tests/test_formatting.py

@@ -392,13 +392,14 @@ def test_torch_formatter_image(self):
         formatter = TorchFormatter(features=Features({"image": Image()}))
         row = formatter.format_row(pa_table)
         self.assertEqual(row["image"].dtype, torch.uint8)
-        self.assertEqual(row["image"].shape, (480, 640, 3))
+        # torch uses CHW format contrary to numpy which uses HWC
+        self.assertEqual(row["image"].shape, (3, 480, 640))
         col = formatter.format_column(pa_table)
         self.assertEqual(col.dtype, torch.uint8)
-        self.assertEqual(col.shape, (2, 480, 640, 3))
+        self.assertEqual(col.shape, (2, 3, 480, 640))
         batch = formatter.format_batch(pa_table)
         self.assertEqual(batch["image"].dtype, torch.uint8)
-        self.assertEqual(batch["image"].shape, (2, 480, 640, 3))
+        self.assertEqual(batch["image"].shape, (2, 3, 480, 640))
 
         # different dimensions
         pa_table = pa.table(
@@ -407,15 +408,15 @@ def test_torch_formatter_image(self):
         formatter = TorchFormatter(features=Features({"image": Image()}))
         row = formatter.format_row(pa_table)
         self.assertEqual(row["image"].dtype, torch.uint8)
-        self.assertEqual(row["image"].shape, (480, 640, 3))
+        self.assertEqual(row["image"].shape, (3, 480, 640))
         col = formatter.format_column(pa_table)
         self.assertIsInstance(col, list)
         self.assertEqual(col[0].dtype, torch.uint8)
-        self.assertEqual(col[0].shape, (480, 640, 3))
+        self.assertEqual(col[0].shape, (3, 480, 640))
         batch = formatter.format_batch(pa_table)
         self.assertIsInstance(batch["image"], list)
         self.assertEqual(batch["image"][0].dtype, torch.uint8)
-        self.assertEqual(batch["image"][0].shape, (480, 640, 3))
+        self.assertEqual(batch["image"][0].shape, (3, 480, 640))
 
     @require_torch
     @require_sndfile