From 16844553f51f11f9223a9f48fc0ad003d3e6ca11 Mon Sep 17 00:00:00 2001 From: Prabhat Roy Date: Thu, 27 May 2021 14:09:06 +0100 Subject: [PATCH 1/2] Fixed missing audio with video_reader backend --- torchvision/io/_video_opt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchvision/io/_video_opt.py b/torchvision/io/_video_opt.py index a34b023bc6c..e92ac1bd396 100644 --- a/torchvision/io/_video_opt.py +++ b/torchvision/io/_video_opt.py @@ -155,7 +155,7 @@ def _align_audio_frames(aframes, aframe_pts, audio_pts_range): e_idx = num_samples if start < audio_pts_range[0]: s_idx = int((audio_pts_range[0] - start) / step_per_aframe) - if end > audio_pts_range[1]: + if audio_pts_range[1] != -1 and end > audio_pts_range[1]: e_idx = int((audio_pts_range[1] - end) / step_per_aframe) return aframes[s_idx:e_idx, :] From e051a415e52e730abb35659e30bd891efd90b249 Mon Sep 17 00:00:00 2001 From: Prabhat Roy Date: Thu, 27 May 2021 17:10:57 +0100 Subject: [PATCH 2/2] Added unit test --- test/test_video_reader.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/test/test_video_reader.py b/test/test_video_reader.py index e13aeb537f2..8f528b429bb 100644 --- a/test/test_video_reader.py +++ b/test/test_video_reader.py @@ -1283,6 +1283,17 @@ def test_invalid_file(self): with self.assertRaises(RuntimeError): io.read_video('foo.mp4') + def test_audio_present(self): + """Test if audio frames are returned with video_reader backend.""" + set_video_backend('video_reader') + for test_video, _ in test_videos.items(): + full_path = os.path.join(VIDEO_DIR, test_video) + container = av.open(full_path) + if container.streams.audio: + _, audio, _ = io.read_video(full_path) + self.assertGreaterEqual(audio.shape[0], 1) + self.assertGreaterEqual(audio.shape[1], 1) + if __name__ == "__main__": unittest.main()