video.py 13.8 KB
Newer Older
1
import gc
2
import math
3
import re
4
import warnings
5
from typing import Any, Dict, List, Optional, Tuple, Union
6

7
8
import numpy as np
import torch
Francisco Massa's avatar
Francisco Massa committed
9

10
11
from . import _video_opt
from ._video_opt import VideoMetaData
Francisco Massa's avatar
Francisco Massa committed
12
13


14
15
try:
    import av
16

17
    av.logging.set_level(av.logging.ERROR)
18
19
20
    if not hasattr(av.video.frame.VideoFrame, "pict_type"):
        av = ImportError(
            """\
21
22
23
24
25
Your version of PyAV is too old for the necessary video operations in torchvision.
If you are on Python 3.5, you will have to build from source (the conda-forge
packages are not up-to-date).  See
https://github.com/mikeboers/PyAV#installation for instructions on how to
install PyAV on your system.
26
27
"""
        )
28
except ImportError:
29
30
    av = ImportError(
        """\
31
32
33
PyAV is not installed, and is necessary for the video operations in torchvision.
See https://github.com/mikeboers/PyAV#installation for instructions on how to
install PyAV on your system.
34
35
"""
    )
36
37


38
def _check_av_available() -> None:
39
40
41
42
    if isinstance(av, Exception):
        raise av


43
def _av_available() -> bool:
44
45
46
    return not isinstance(av, Exception)


47
48
# PyAV has some reference cycles
_CALLED_TIMES = 0
49
_GC_COLLECTION_INTERVAL = 10
50
51


52
53
54
55
56
57
def write_video(
    filename: str,
    video_array: torch.Tensor,
    fps: float,
    video_codec: str = "libx264",
    options: Optional[Dict[str, Any]] = None,
58
59
60
61
    audio_array: Optional[torch.Tensor] = None,
    audio_fps: Optional[float] = None,
    audio_codec: Optional[str] = None,
    audio_options: Optional[Dict[str, Any]] = None,
62
) -> None:
63
64
65
    """
    Writes a 4d tensor in [T, H, W, C] format in a video file

66
67
68
69
70
71
72
    Parameters
    ----------
    filename : str
        path where the video will be saved
    video_array : Tensor[T, H, W, C]
        tensor containing the individual frames, as a uint8 tensor in [T, H, W, C] format
    fps : Number
73
74
75
76
77
78
79
80
81
82
83
84
85
86
        video frames per second
    video_codec : str
        the name of the video codec, i.e. "libx264", "h264", etc.
    options : Dict
        dictionary containing options to be passed into the PyAV video stream
    audio_array : Tensor[C, N]
        tensor containing the audio, where C is the number of channels and N is the
        number of samples
    audio_fps : Number
        audio sample rate, typically 44100 or 48000
    audio_codec : str
        the name of the audio codec, i.e. "mp3", "aac", etc.
    audio_options : Dict
        dictionary containing options to be passed into the PyAV audio stream
87
88
89
90
    """
    _check_av_available()
    video_array = torch.as_tensor(video_array, dtype=torch.uint8).numpy()

91
92
93
94
95
    # PyAV does not support floating point numbers with decimal point
    # and will throw OverflowException in case this is not the case
    if isinstance(fps, float):
        fps = np.round(fps)

96
97
98
99
100
101
102
    with av.open(filename, mode="w") as container:
        stream = container.add_stream(video_codec, rate=fps)
        stream.width = video_array.shape[2]
        stream.height = video_array.shape[1]
        stream.pix_fmt = "yuv420p" if video_codec != "libx264rgb" else "rgb24"
        stream.options = options or {}

103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
        if audio_array is not None:
            audio_format_dtypes = {
                'dbl': '<f8',
                'dblp': '<f8',
                'flt': '<f4',
                'fltp': '<f4',
                's16': '<i2',
                's16p': '<i2',
                's32': '<i4',
                's32p': '<i4',
                'u8': 'u1',
                'u8p': 'u1',
            }
            a_stream = container.add_stream(audio_codec, rate=audio_fps)
            a_stream.options = audio_options or {}

            num_channels = audio_array.shape[0]
            audio_layout = "stereo" if num_channels > 1 else "mono"
            audio_sample_fmt = container.streams.audio[0].format.name

            format_dtype = np.dtype(audio_format_dtypes[audio_sample_fmt])
            audio_array = torch.as_tensor(audio_array).numpy().astype(format_dtype)

            frame = av.AudioFrame.from_ndarray(
                audio_array, format=audio_sample_fmt, layout=audio_layout
            )

            frame.sample_rate = audio_fps

            for packet in a_stream.encode(frame):
                container.mux(packet)

            for packet in a_stream.encode():
                container.mux(packet)

138
139
140
141
142
143
144
145
        for img in video_array:
            frame = av.VideoFrame.from_ndarray(img, format="rgb24")
            frame.pict_type = "NONE"
            for packet in stream.encode(frame):
                container.mux(packet)

        # Flush stream
        for packet in stream.encode():
146
147
148
            container.mux(packet)


149
def _read_from_stream(
150
151
152
153
154
155
156
    container: "av.container.Container",
    start_offset: float,
    end_offset: float,
    pts_unit: str,
    stream: "av.stream.Stream",
    stream_name: Dict[str, Optional[Union[int, Tuple[int, ...], List[int]]]],
) -> List["av.frame.Frame"]:
157
158
159
160
161
    global _CALLED_TIMES, _GC_COLLECTION_INTERVAL
    _CALLED_TIMES += 1
    if _CALLED_TIMES % _GC_COLLECTION_INTERVAL == _GC_COLLECTION_INTERVAL - 1:
        gc.collect()

162
    if pts_unit == "sec":
163
164
165
166
        start_offset = int(math.floor(start_offset * (1 / stream.time_base)))
        if end_offset != float("inf"):
            end_offset = int(math.ceil(end_offset * (1 / stream.time_base)))
    else:
167
168
169
170
        warnings.warn(
            "The pts_unit 'pts' gives wrong results and will be removed in a "
            + "follow-up version. Please use pts_unit 'sec'."
        )
171

172
    frames = {}
173
    should_buffer = True
174
175
    max_buffer_size = 5
    if stream.type == "video":
176
        # DivX-style packed B-frames can have out-of-order pts (2 frames in a single pkt)
177
178
        # so need to buffer some extra frames to sort everything
        # properly
179
180
181
182
183
184
185
186
187
188
189
190
        extradata = stream.codec_context.extradata
        # overly complicated way of finding if `divx_packed` is set, following
        # https://github.com/FFmpeg/FFmpeg/commit/d5a21172283572af587b3d939eba0091484d3263
        if extradata and b"DivX" in extradata:
            # can't use regex directly because of some weird characters sometimes...
            pos = extradata.find(b"DivX")
            d = extradata[pos:]
            o = re.search(br"DivX(\d+)Build(\d+)(\w)", d)
            if o is None:
                o = re.search(br"DivX(\d+)b(\d+)(\w)", d)
            if o is not None:
                should_buffer = o.group(3) == b"p"
191
    seek_offset = start_offset
192
193
    # some files don't seek to the right location, so better be safe here
    seek_offset = max(seek_offset - 1, 0)
194
195
196
197
    if should_buffer:
        # FIXME this is kind of a hack, but we will jump to the previous keyframe
        # so this will be safe
        seek_offset = max(seek_offset - max_buffer_size, 0)
198
199
200
201
    try:
        # TODO check if stream needs to always be the video stream here or not
        container.seek(seek_offset, any_frame=False, backward=True, stream=stream)
    except av.AVError:
202
203
        # TODO add some warnings in this case
        # print("Corrupted file?", container.name)
204
        return []
205
    buffer_count = 0
206
    try:
207
        for _idx, frame in enumerate(container.decode(**stream_name)):
208
209
210
211
212
213
214
215
216
            frames[frame.pts] = frame
            if frame.pts >= end_offset:
                if should_buffer and buffer_count < max_buffer_size:
                    buffer_count += 1
                    continue
                break
    except av.AVError:
        # TODO add a warning
        pass
217
    # ensure that the results are sorted wrt the pts
218
219
220
    result = [
        frames[i] for i in sorted(frames) if start_offset <= frames[i].pts <= end_offset
    ]
221
    if len(frames) > 0 and start_offset > 0 and start_offset not in frames:
222
223
224
        # if there is no frame that exactly matches the pts of start_offset
        # add the last frame smaller than start_offset, to guarantee that
        # we will have all the necessary data. This is most useful for audio
225
226
227
228
        preceding_frames = [i for i in frames if i < start_offset]
        if len(preceding_frames) > 0:
            first_frame_pts = max(preceding_frames)
            result.insert(0, frames[first_frame_pts])
229
    return result
230
231


232
233
234
def _align_audio_frames(
    aframes: torch.Tensor, audio_frames: List["av.frame.Frame"], ref_start: int, ref_end: float
) -> torch.Tensor:
235
236
237
238
239
240
241
242
243
244
245
246
    start, end = audio_frames[0].pts, audio_frames[-1].pts
    total_aframes = aframes.shape[1]
    step_per_aframe = (end - start + 1) / total_aframes
    s_idx = 0
    e_idx = total_aframes
    if start < ref_start:
        s_idx = int((ref_start - start) / step_per_aframe)
    if end > ref_end:
        e_idx = int((ref_end - end) / step_per_aframe)
    return aframes[:, s_idx:e_idx]


247
248
249
def read_video(
    filename: str, start_pts: int = 0, end_pts: Optional[float] = None, pts_unit: str = "pts"
) -> Tuple[torch.Tensor, torch.Tensor, Dict[str, Any]]:
250
251
252
253
    """
    Reads a video from a file, returning both the video frames as well as
    the audio frames

254
255
256
257
    Parameters
    ----------
    filename : str
        path to the video file
258
259
    start_pts : int if pts_unit = 'pts', optional
        float / Fraction if pts_unit = 'sec', optional
260
        the start presentation time of the video
261
262
    end_pts : int if pts_unit = 'pts', optional
        float / Fraction if pts_unit = 'sec', optional
263
        the end presentation time
264
265
    pts_unit : str, optional
        unit in which start_pts and end_pts values will be interpreted, either 'pts' or 'sec'. Defaults to 'pts'.
266
267
268
269
270
271
272
273
274
275
276

    Returns
    -------
    vframes : Tensor[T, H, W, C]
        the `T` video frames
    aframes : Tensor[K, L]
        the audio frames, where `K` is the number of channels and `L` is the
        number of points
    info : Dict
        metadata for the video and audio. Can contain the fields video_fps (float)
        and audio_fps (int)
277
    """
Francisco Massa's avatar
Francisco Massa committed
278
279

    from torchvision import get_video_backend
280

Francisco Massa's avatar
Francisco Massa committed
281
282
283
    if get_video_backend() != "pyav":
        return _video_opt._read_video(filename, start_pts, end_pts, pts_unit)

284
285
286
287
288
289
    _check_av_available()

    if end_pts is None:
        end_pts = float("inf")

    if end_pts < start_pts:
290
291
292
293
        raise ValueError(
            "end_pts should be larger than start_pts, got "
            "start_pts={} and end_pts={}".format(start_pts, end_pts)
        )
294
295
296
297
298

    info = {}
    video_frames = []
    audio_frames = []

299
    try:
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
        with av.open(filename, metadata_errors="ignore") as container:
            if container.streams.video:
                video_frames = _read_from_stream(
                    container,
                    start_pts,
                    end_pts,
                    pts_unit,
                    container.streams.video[0],
                    {"video": 0},
                )
                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,
                    container.streams.audio[0],
                    {"audio": 0},
                )
                info["audio_fps"] = container.streams.audio[0].rate

326
327
328
    except av.AVError:
        # TODO raise a warning?
        pass
329

330
331
    vframes_list = [frame.to_rgb().to_ndarray() for frame in video_frames]
    aframes_list = [frame.to_ndarray() for frame in audio_frames]
332

333
334
    if vframes_list:
        vframes = torch.as_tensor(np.stack(vframes_list))
335
336
337
    else:
        vframes = torch.empty((0, 1, 1, 3), dtype=torch.uint8)

338
339
    if aframes_list:
        aframes = np.concatenate(aframes_list, 1)
340
341
342
343
344
345
346
347
        aframes = torch.as_tensor(aframes)
        aframes = _align_audio_frames(aframes, audio_frames, start_pts, end_pts)
    else:
        aframes = torch.empty((1, 0), dtype=torch.float32)

    return vframes, aframes, info


348
def _can_read_timestamps_from_packets(container: "av.container.Container") -> bool:
349
350
351
352
353
354
355
356
    extradata = container.streams[0].codec_context.extradata
    if extradata is None:
        return False
    if b"Lavc" in extradata:
        return True
    return False


357
def _decode_video_timestamps(container: "av.container.Container") -> List[int]:
358
359
360
361
362
363
364
    if _can_read_timestamps_from_packets(container):
        # fast path
        return [x.pts for x in container.demux(video=0) if x.pts is not None]
    else:
        return [x.pts for x in container.decode(video=0) if x.pts is not None]


365
def read_video_timestamps(filename: str, pts_unit: str = "pts") -> Tuple[List[int], Optional[float]]:
366
367
368
369
370
    """
    List the video frames timestamps.

    Note that the function decodes the whole video frame-by-frame.

371
372
373
374
    Parameters
    ----------
    filename : str
        path to the video file
375
376
    pts_unit : str, optional
        unit in which timestamp values will be returned either 'pts' or 'sec'. Defaults to 'pts'.
377
378
379

    Returns
    -------
380
381
    pts : List[int] if pts_unit = 'pts'
        List[Fraction] if pts_unit = 'sec'
382
        presentation timestamps for each one of the frames in the video.
383
    video_fps : float, optional
384
        the frame rate for the video
385
386

    """
Francisco Massa's avatar
Francisco Massa committed
387
    from torchvision import get_video_backend
388

Francisco Massa's avatar
Francisco Massa committed
389
390
391
    if get_video_backend() != "pyav":
        return _video_opt._read_video_timestamps(filename, pts_unit)

392
    _check_av_available()
393

394
    video_fps = None
395
    pts = []
396
397

    try:
398
399
400
401
402
403
404
405
406
        with av.open(filename, metadata_errors="ignore") as container:
            if container.streams.video:
                video_stream = container.streams.video[0]
                video_time_base = video_stream.time_base
                try:
                    pts = _decode_video_timestamps(container)
                except av.AVError:
                    warnings.warn(f"Failed decoding frames for file {filename}")
                video_fps = float(video_stream.average_rate)
407
408
409
410
    except av.AVError:
        # TODO add a warning
        pass

411
    pts.sort()
412

413
    if pts_unit == "sec":
414
415
416
        pts = [x * video_time_base for x in pts]

    return pts, video_fps