[fbsync] Fixed missing audio with video_reader backend (#3934)

Summary:
* Fixed missing audio with video_reader backend

* Added unit test

Reviewed By: fmassa

Differential Revision: D29264318

fbshipit-source-id: de95e0bd38d2f844c756652fe42de99b1ab32210

Author: NicolasHug

test/test_video_reader.py

@@ -1244,6 +1244,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()

torchvision/io/_video_opt.py

@@ -122,7 +122,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, :]