Added tests for torch_to_matplotlib in visualization.py

Signed-off-by: Martin <[email protected]>

dgs/utils/visualization.py

@@ -46,7 +46,7 @@ def torch_show_image(
         plt.show()
 
 
-def torch_to_matplotlib(img: TVImage) -> np.ndarray:  # pragma: no cover
+def torch_to_matplotlib(img: Union[TVImage, torch.Tensor]) -> np.ndarray:
     """Convert a given single or batched torch image Tensor to a numpy.ndarray on the cpu.
     The dimensions are switched from ``[B x C x H x W]`` -> ``[B x H x W x C]``
 

tests/utils/test__visualization.py

@@ -0,0 +1,23 @@
+import unittest
+
+import torch
+from torchvision import tv_tensors
+
+from dgs.utils.visualization import torch_to_matplotlib
+
+
+class TestVisualization(unittest.TestCase):
+    def test_torch_to_matplotlib(self):
+        B, C, H, W = 8, 3, 64, 64
+        for tensor, out_shape in [
+            (torch.ones(B, C, H, W), [B, H, W, C]),
+            (tv_tensors.Image(torch.ones(B, C, H, W)), [B, H, W, C]),
+            (torch.ones(C, H, W), [H, W, C]),
+        ]:
+            with self.subTest(msg="tensor: {}, out_shape: {}".format(tensor, out_shape)):
+                m = torch_to_matplotlib(tensor)
+                self.assertEqual(list(m.shape), list(out_shape))
+
+
+if __name__ == "__main__":
+    unittest.main()