diff --git a/torchvision/io/video.py b/torchvision/io/video.py index bd25c224ecb..7b340cd9766 100644 --- a/torchvision/io/video.py +++ b/torchvision/io/video.py @@ -2,6 +2,8 @@ import gc import torch import numpy as np +import math +import warnings try: import av @@ -145,7 +147,7 @@ def _align_audio_frames(aframes, audio_frames, ref_start, ref_end): return aframes[:, s_idx:e_idx] -def read_video(filename, start_pts=0, end_pts=None): +def read_video(filename, start_pts=0, end_pts=None, pts_unit='sec'): """ Reads a video from a file, returning both the video frames as well as the audio frames @@ -158,6 +160,8 @@ def read_video(filename, start_pts=0, end_pts=None): the start presentation time of the video end_pts : int, optional the end presentation time + pts_unit : str, optional + unit in which start_pts and end_pts values will be interpreted, either 'pts' or 'sec. Defaults to 'sec'. Returns ------- @@ -184,14 +188,34 @@ def read_video(filename, start_pts=0, end_pts=None): video_frames = [] if container.streams.video: - video_frames = _read_from_stream(container, start_pts, end_pts, - container.streams.video[0], {'video': 0}) - info["video_fps"] = float(container.streams.video[0].average_rate) + video_stream = container.streams.video[0] + if pts_unit == 'pts': + warnings.warn("The pts_unit 'pts' produces wrong results and will be removed + in a follow-up version. Please use pts_unit 'sec'.") + video_start_pts = start_pts + video_end_pts = end_pts + else: + video_start_pts = math.floor(start_pts*(1/video_stream.time_base)) + if video_end_pts != float("inf"): + video_end_pts = math.ceil(end_pts*(1/video_stream.time_base)) + video_frames = _read_from_stream(container, video_start_pts, video_end_pts, + video_stream, {'video': 0}) + info["video_fps"] = float(video_stream.average_rate) audio_frames = [] if container.streams.audio: - audio_frames = _read_from_stream(container, start_pts, end_pts, - container.streams.audio[0], {'audio': 0}) - info["audio_fps"] = container.streams.audio[0].rate + audio_stream = container.streams.audio[0] + if pts_unit == 'pts': + warnings.warn("The pts_unit 'pts' produces wrong results and will be removed + in a follow-up version. Please use pts_unit 'sec'.") + audio_start_pts = start_pts + audio_end_pts = end_pts + else: + audio_start_pts = math.floor(start_pts*(1/audio_stream.time_base)) + if end_pts != float("inf"): + audio_end_pts = math.ceil(end_pts*(1/audio_stream.time_base)) + audio_frames = _read_from_stream(container, audio_start_pts, audio_end_pts, + audio_stream , {'audio': 0}) + info["audio_fps"] = audio_stream.rate container.close() @@ -217,7 +241,7 @@ def _can_read_timestamps_from_packets(container): return False -def read_video_timestamps(filename): +def read_video_timestamps(filename, pts_unit='pts'): """ List the video frames timestamps. @@ -227,6 +251,8 @@ def read_video_timestamps(filename): ---------- filename : str path to the video file + pts_unit : str, optional + unit in which timestamp values will be returned either 'pts' or 'sec. Defaults to 'sec'. Returns ------- @@ -242,12 +268,18 @@ def read_video_timestamps(filename): video_frames = [] video_fps = None if container.streams.video: + video_stream = container.streams.video[0] if _can_read_timestamps_from_packets(container): # fast path video_frames = [x for x in container.demux(video=0) if x.pts is not None] else: video_frames = _read_from_stream(container, 0, float("inf"), - container.streams.video[0], {'video': 0}) - video_fps = float(container.streams.video[0].average_rate) + video_stream, {'video': 0}) + video_fps = float(video_stream.average_rate) container.close() - return [x.pts for x in video_frames], video_fps + if pts_unit == 'pts': + warnings.warn("The pts_unit 'pts' produces wrong results and will be removed in a + follow-up version. Please use pts_unit 'sec'.") + return [x.pts for x in video_frames], video_fps + else: + return [x.pts*video_stream.time_base for x in video_frames], video_fps