add test_make_batched_videos

Author: yonigozlan

src/transformers/image_utils.py

@@ -292,9 +292,6 @@ def make_nested_list_of_images(
             return [list(image) for image in images]
 
     # If it's a single image, convert it to a list of lists
-    if is_pil_image(images):
-        return [[images]]
-
     if is_valid_image(images):
         if is_pil_image(images) or images.ndim == 3:
             return [[images]]
@@ -317,15 +314,15 @@ def make_batched_videos(videos) -> VideoInput:
         return videos
 
     elif isinstance(videos, (list, tuple)) and is_valid_image(videos[0]):
-        if is_pil_image(videos[0]):
+        if is_pil_image(videos[0]) or videos[0].ndim == 3:
             return [videos]
-        elif len(videos[0].shape) == 4:
+        elif videos[0].ndim == 4:
             return [list(video) for video in videos]
 
     elif is_valid_image(videos):
-        if is_pil_image(videos):
+        if is_pil_image(videos) or videos.ndim == 3:
             return [[videos]]
-        elif len(videos.shape) == 4:
+        elif videos.ndim == 4:
             return [list(videos)]
 
     raise ValueError(f"Could not make batched video from {videos}")

tests/utils/test_image_utils.py

@@ -31,6 +31,7 @@
 from transformers.image_utils import (
     ChannelDimension,
     get_channel_dimension_axis,
+    make_batched_videos,
     make_flat_list_of_images,
     make_list_of_images,
     make_nested_list_of_images,
@@ -396,6 +397,115 @@ def test_make_nested_list_of_images_torch(self):
         self.assertEqual(len(images_list[0]), 4)
         self.assertTrue(np.array_equal(images_list[0][0], images[0][0]))
 
+    def test_make_batched_videos_pil(self):
+        # Test a single image is converted to a list of 1 video with 1 frame
+        pil_image = get_random_image(16, 32)
+        videos_list = make_batched_videos(pil_image)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list[0]), 1)
+        self.assertIsInstance(videos_list[0][0], PIL.Image.Image)
+
+        # Test a list of images is converted to a list of 1 video
+        images = [get_random_image(16, 32) for _ in range(4)]
+        videos_list = make_batched_videos(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertIsInstance(videos_list[0][0], PIL.Image.Image)
+
+        # Test a nested list of images is not modified
+        images = [[get_random_image(16, 32) for _ in range(2)] for _ in range(2)]
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 2)
+        self.assertIsInstance(videos_list[0][0], PIL.Image.Image)
+
+    def test_make_batched_videos_numpy(self):
+        # Test a single image is converted to a list of 1 video with 1 frame
+        images = np.random.randint(0, 256, (16, 32, 3))
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 1)
+        self.assertTrue(np.array_equal(videos_list[0][0], images))
+
+        # Test a 4d array of images is converted to a a list of 1 video
+        images = np.random.randint(0, 256, (4, 16, 32, 3))
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], np.ndarray)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0]))
+
+        # Test a list of images is converted to a list of videos
+        images = [np.random.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0]))
+
+        # Test a nested list of images is left unchanged
+        images = [[np.random.randint(0, 256, (16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 2)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0]))
+
+        # Test a list of 4d array images is converted to a list of videos
+        images = [np.random.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)]
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], np.ndarray)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0]))
+
+    @require_torch
+    def test_make_batched_videos_torch(self):
+        # Test a single image is converted to a list of 1 video with 1 frame
+        images = torch.randint(0, 256, (16, 32, 3))
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list[0]), 1)
+        self.assertTrue(np.array_equal(videos_list[0][0], images))
+
+        # Test a 4d tensor of images is converted to a list of 1 video
+        images = torch.randint(0, 256, (4, 16, 32, 3))
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], torch.Tensor)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0]))
+
+        # Test a list of images is converted to a list of videos
+        images = [torch.randint(0, 256, (16, 32, 3)) for _ in range(4)]
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 1)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0]))
+
+        # Test a nested list of images is left unchanged
+        images = [[torch.randint(0, 256, (16, 32, 3)) for _ in range(2)] for _ in range(2)]
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 2)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0]))
+
+        # Test a list of 4d tensor images is converted to a list of videos
+        images = [torch.randint(0, 256, (4, 16, 32, 3)) for _ in range(2)]
+        videos_list = make_nested_list_of_images(images)
+        self.assertIsInstance(videos_list[0], list)
+        self.assertIsInstance(videos_list[0][0], torch.Tensor)
+        self.assertEqual(len(videos_list), 2)
+        self.assertEqual(len(videos_list[0]), 4)
+        self.assertTrue(np.array_equal(videos_list[0][0], images[0][0]))
+
     @require_torch
     def test_conversion_torch_to_array(self):
         feature_extractor = ImageFeatureExtractionMixin()