Change error in combine_images_to_video and add test

Signed-off-by: Martin <[email protected]>
bmmtstb · Apr 6, 2024 · ffad5df · ffad5df
1 parent 034395c
commit ffad5df
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/dgs/utils/image.py b/dgs/utils/image.py
@@ -189,7 +189,7 @@ def combine_images_to_video(
     elif isinstance(imgs, list):
         images = transform_dtype(torch.cat(imgs))
     else:
-        raise NotImplementedError(f"Unknown input format.")
+        raise TypeError(f"Unknown input format. Got {type(imgs)}")
 
     # change order of the dimensions
     video_tensor = torch.permute(images, (0, 2, 3, 1))

diff --git a/tests/utils/test__image.py b/tests/utils/test__image.py
@@ -176,6 +176,11 @@ def test_combine_images_to_video(self):
 
         shutil.rmtree("./tests/test_data/video_out/")
 
+    def test_combine_exceptions(self):
+        with self.assertRaises(TypeError) as e:
+            combine_images_to_video(imgs=np.ones(1), video_file="./tests/test_data/video_out/test1.mp4")
+        self.assertTrue("Unknown input format. Got" in str(e.exception), msg=e.exception)
+
 
 class TestImage(unittest.TestCase):