Add more tests

Author: fmassa

test/test_io.py

@@ -253,6 +253,24 @@ def test_read_video_timestamps_corrupted_file(self):
             self.assertEqual(video_pts, [])
             self.assertIs(video_fps, None)
 
+    def test_read_video_partially_corrupted_file(self):
+        with temp_video(5, 4, 4, 5, lossless=True) as (f_name, data):
+            with open(f_name, 'r+b') as f:
+                size = os.path.getsize(f_name)
+                bytes_to_overwrite = size // 10
+                # seek to the middle of the file
+                f.seek(5 * bytes_to_overwrite)
+                # corrupt 10% of the file from the middle
+                f.write(b'\xff' * bytes_to_overwrite)
+            # this exercises the container.decode assertion check
+            video, audio, info = io.read_video(f.name, pts_unit='sec')
+            # check that size is not equal to 5, but 3
+            self.assertEqual(len(video), 3)
+            # but the valid decoded content is still correct
+            self.assertTrue(video[:3].equal(data[:3]))
+            # and the last few frames are wrong
+            self.assertFalse(video.equal(data))
+
     # TODO add tests for audio
 
 

torchvision/io/video.py

@@ -210,7 +210,10 @@ def read_video(filename, start_pts=0, end_pts=None, pts_unit='pts'):
         if container.streams.video:
             video_frames = _read_from_stream(container, start_pts, end_pts, pts_unit,
                                              container.streams.video[0], {'video': 0})
-            info["video_fps"] = float(container.streams.video[0].average_rate)
+            video_fps = container.streams.video[0].average_rate
+            # guard against potentially corrupted files
+            if video_fps is not None:
+                info["video_fps"] = float(video_fps)
 
         if container.streams.audio:
             audio_frames = _read_from_stream(container, start_pts, end_pts, pts_unit,