Skip to content

Commit

Permalink
Fixed missing audio with video_reader backend (#3934)
Browse files Browse the repository at this point in the history
* Fixed missing audio with video_reader backend

* Added unit test
  • Loading branch information
prabhat00155 authored Jun 15, 2021
1 parent 496cb40 commit b74366b
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 1 deletion.
11 changes: 11 additions & 0 deletions test/test_video_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
2 changes: 1 addition & 1 deletion torchvision/io/_video_opt.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, :]

Expand Down

0 comments on commit b74366b

Please sign in to comment.