video_utils.py 15.1 KB
Newer Older
1
2
import bisect
import math
3
import warnings
4
from fractions import Fraction
5
from typing import List
6

7
import torch
8
from torchvision.io import (
9
    _probe_video_from_file,
10
    _read_video_from_file,
11
12
    read_video,
    read_video_timestamps,
13
)
14

15
16
from .utils import tqdm

17

18
19
20
21
22
23
24
25
26
27
28
29
def pts_convert(pts, timebase_from, timebase_to, round_func=math.floor):
    """convert pts between different time bases
    Args:
        pts: presentation timestamp, float
        timebase_from: original timebase. Fraction
        timebase_to: new timebase. Fraction
        round_func: rounding function.
    """
    new_pts = Fraction(pts, 1) * timebase_from / timebase_to
    return round_func(new_pts)


30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def unfold(tensor, size, step, dilation=1):
    """
    similar to tensor.unfold, but with the dilation
    and specialized for 1d tensors

    Returns all consecutive windows of `size` elements, with
    `step` between windows. The distance between each element
    in a window is given by `dilation`.
    """
    assert tensor.dim() == 1
    o_stride = tensor.stride(0)
    numel = tensor.numel()
    new_stride = (step * o_stride, dilation * o_stride)
    new_size = ((numel - (dilation * (size - 1) + 1)) // step + 1, size)
    if new_size[0] < 1:
        new_size = (0, size)
    return torch.as_strided(tensor, new_size, new_stride)


49
class _VideoTimestampsDataset:
50
    """
51
52
53
54
55
    Dataset used to parallelize the reading of the timestamps
    of a list of videos, given their paths in the filesystem.

    Used in VideoClips and defined at top level so it can be
    pickled when forking.
56
    """
57

58
59
    def __init__(self, video_paths: List[str]):
        self.video_paths = video_paths
60
61

    def __len__(self):
62
        return len(self.video_paths)
63
64

    def __getitem__(self, idx):
65
        return read_video_timestamps(self.video_paths[idx])
66
67


68
69
70
71
72
73
74
def _collate_fn(x):
    """
    Dummy collate function to be used with _VideoTimestampsDataset
    """
    return x


75
class VideoClips:
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    """
    Given a list of video files, computes all consecutive subvideos of size
    `clip_length_in_frames`, where the distance between each subvideo in the
    same video is defined by `frames_between_clips`.
    If `frame_rate` is specified, it will also resample all the videos to have
    the same frame rate, and the clips will refer to this frame rate.

    Creating this instance the first time is time-consuming, as it needs to
    decode all the videos in `video_paths`. It is recommended that you
    cache the results after instantiation of the class.

    Recreating the clips for different clip lengths is fast, and can be done
    with the `compute_clips` method.

90
    Args:
91
92
93
94
95
96
        video_paths (List[str]): paths to the video files
        clip_length_in_frames (int): size of a clip in number of frames
        frames_between_clips (int): step (in frames) between each clip
        frame_rate (int, optional): if specified, it will resample the video
            so that it has `frame_rate`, and then the clips will be defined
            on the resampled video
ekosman's avatar
ekosman committed
97
98
        num_workers (int): how many subprocesses to use for data loading.
            0 means that the data will be loaded in the main process. (default: 0)
99
    """
100
101
102
103
104
105
106
107
108
109
110
111

    def __init__(
        self,
        video_paths,
        clip_length_in_frames=16,
        frames_between_clips=1,
        frame_rate=None,
        _precomputed_metadata=None,
        num_workers=0,
        _video_width=0,
        _video_height=0,
        _video_min_dimension=0,
112
        _video_max_dimension=0,
113
114
115
        _audio_samples=0,
        _audio_channels=0,
    ):
116

117
        self.video_paths = video_paths
118
        self.num_workers = num_workers
119
120

        # these options are not valid for pyav backend
121
122
123
        self._video_width = _video_width
        self._video_height = _video_height
        self._video_min_dimension = _video_min_dimension
124
        self._video_max_dimension = _video_max_dimension
125
        self._audio_samples = _audio_samples
126
        self._audio_channels = _audio_channels
ekosman's avatar
ekosman committed
127

128
129
130
131
        if _precomputed_metadata is None:
            self._compute_frame_pts()
        else:
            self._init_from_metadata(_precomputed_metadata)
132
133
134
135
        self.compute_clips(clip_length_in_frames, frames_between_clips, frame_rate)

    def _compute_frame_pts(self):
        self.video_pts = []
136
        self.video_fps = []
137
138
139
140

        # strategy: use a DataLoader to parallelize read_video_timestamps
        # so need to create a dummy dataset first
        import torch.utils.data
141

142
        dl = torch.utils.data.DataLoader(
143
            _VideoTimestampsDataset(self.video_paths),
144
            batch_size=16,
145
            num_workers=self.num_workers,
146
            collate_fn=_collate_fn,
147
        )
148
149
150
151

        with tqdm(total=len(dl)) as pbar:
            for batch in dl:
                pbar.update(1)
152
                clips, fps = list(zip(*batch))
153
154
155
156
                # we need to specify dtype=torch.long because for empty list,
                # torch.as_tensor will use torch.float as default dtype. This
                # happens when decoding fails and no pts is returned in the list.
                clips = [torch.as_tensor(c, dtype=torch.long) for c in clips]
157
158
                self.video_pts.extend(clips)
                self.video_fps.extend(fps)
159

160
    def _init_from_metadata(self, metadata):
161
        self.video_paths = metadata["video_paths"]
162
163
        assert len(self.video_paths) == len(metadata["video_pts"])
        self.video_pts = metadata["video_pts"]
164
165
        assert len(self.video_paths) == len(metadata["video_fps"])
        self.video_fps = metadata["video_fps"]
166
167
168
169
170
171

    @property
    def metadata(self):
        _metadata = {
            "video_paths": self.video_paths,
            "video_pts": self.video_pts,
172
            "video_fps": self.video_fps,
173
        }
174
        return _metadata
175
176
177
178

    def subset(self, indices):
        video_paths = [self.video_paths[i] for i in indices]
        video_pts = [self.video_pts[i] for i in indices]
179
        video_fps = [self.video_fps[i] for i in indices]
180
        metadata = {
181
            "video_paths": video_paths,
182
            "video_pts": video_pts,
183
            "video_fps": video_fps,
184
        }
185
186
187
188
189
190
191
192
193
194
        return type(self)(
            video_paths,
            self.num_frames,
            self.step,
            self.frame_rate,
            _precomputed_metadata=metadata,
            num_workers=self.num_workers,
            _video_width=self._video_width,
            _video_height=self._video_height,
            _video_min_dimension=self._video_min_dimension,
195
            _video_max_dimension=self._video_max_dimension,
196
197
198
            _audio_samples=self._audio_samples,
            _audio_channels=self._audio_channels,
        )
199

200
201
    @staticmethod
    def compute_clips_for_video(video_pts, num_frames, step, fps, frame_rate):
202
203
204
205
        if fps is None:
            # if for some reason the video doesn't have fps (because doesn't have a video stream)
            # set the fps to 1. The value doesn't matter, because video_pts is empty anyway
            fps = 1
206
207
208
        if frame_rate is None:
            frame_rate = fps
        total_frames = len(video_pts) * (float(frame_rate) / fps)
209
        idxs = VideoClips._resample_video_idx(int(math.floor(total_frames)), fps, frame_rate)
210
211
        video_pts = video_pts[idxs]
        clips = unfold(video_pts, num_frames, step)
212
        if not clips.numel():
213
214
215
216
            warnings.warn(
                "There aren't enough frames in the current video to get a clip for the given clip length and "
                "frames between clips. The video (and potentially others) will be skipped."
            )
217
218
219
220
221
222
223
224
225
226
227
228
        if isinstance(idxs, slice):
            idxs = [idxs] * len(clips)
        else:
            idxs = unfold(idxs, num_frames, step)
        return clips, idxs

    def compute_clips(self, num_frames, step, frame_rate=None):
        """
        Compute all consecutive sequences of clips from video_pts.
        Always returns clips of size `num_frames`, meaning that the
        last few frames in a video can potentially be dropped.

229
        Args:
230
231
            num_frames (int): number of frames for the clip
            step (int): distance between two clips
232
            frame_rate (int, optional): The frame rate
233
234
235
236
237
238
        """
        self.num_frames = num_frames
        self.step = step
        self.frame_rate = frame_rate
        self.clips = []
        self.resampling_idxs = []
239
        for video_pts, fps in zip(self.video_pts, self.video_fps):
240
            clips, idxs = self.compute_clips_for_video(video_pts, num_frames, step, fps, frame_rate)
241
242
            self.clips.append(clips)
            self.resampling_idxs.append(idxs)
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
        clip_lengths = torch.as_tensor([len(v) for v in self.clips])
        self.cumulative_sizes = clip_lengths.cumsum(0).tolist()

    def __len__(self):
        return self.num_clips()

    def num_videos(self):
        return len(self.video_paths)

    def num_clips(self):
        """
        Number of subclips that are available in the video list.
        """
        return self.cumulative_sizes[-1]

    def get_clip_location(self, idx):
        """
        Converts a flattened representation of the indices into a video_idx, clip_idx
        representation.
        """
        video_idx = bisect.bisect_right(self.cumulative_sizes, idx)
        if video_idx == 0:
            clip_idx = idx
        else:
            clip_idx = idx - self.cumulative_sizes[video_idx - 1]
        return video_idx, clip_idx

    @staticmethod
    def _resample_video_idx(num_frames, original_fps, new_fps):
        step = float(original_fps) / new_fps
        if step.is_integer():
            # optimization: if step is integer, don't need to perform
            # advanced indexing
            step = int(step)
            return slice(None, None, step)
        idxs = torch.arange(num_frames, dtype=torch.float32) * step
        idxs = idxs.floor().to(torch.int64)
        return idxs

    def get_clip(self, idx):
        """
        Gets a subclip from a list of videos.

286
        Args:
287
288
289
290
291
292
293
294
295
            idx (int): index of the subclip. Must be between 0 and num_clips().

        Returns:
            video (Tensor)
            audio (Tensor)
            info (Dict)
            video_idx (int): index of the video in `video_paths`
        """
        if idx >= self.num_clips():
296
            raise IndexError(f"Index {idx} out of range ({self.num_clips()} number of clips)")
297
298
299
        video_idx, clip_idx = self.get_clip_location(idx)
        video_path = self.video_paths[video_idx]
        clip_pts = self.clips[video_idx][clip_idx]
300

301
        from torchvision import get_video_backend
302

303
304
305
306
307
308
309
310
311
        backend = get_video_backend()

        if backend == "pyav":
            # check for invalid options
            if self._video_width != 0:
                raise ValueError("pyav backend doesn't support _video_width != 0")
            if self._video_height != 0:
                raise ValueError("pyav backend doesn't support _video_height != 0")
            if self._video_min_dimension != 0:
312
                raise ValueError("pyav backend doesn't support _video_min_dimension != 0")
313
            if self._video_max_dimension != 0:
314
                raise ValueError("pyav backend doesn't support _video_max_dimension != 0")
315
316
317
318
            if self._audio_samples != 0:
                raise ValueError("pyav backend doesn't support _audio_samples != 0")

        if backend == "pyav":
319
320
321
322
            start_pts = clip_pts[0].item()
            end_pts = clip_pts[-1].item()
            video, audio, info = read_video(video_path, start_pts, end_pts)
        else:
323
            info = _probe_video_from_file(video_path)
324
            video_fps = info.video_fps
325
            audio_fps = None
326
327
328
329
330
331

            video_start_pts = clip_pts[0].item()
            video_end_pts = clip_pts[-1].item()

            audio_start_pts, audio_end_pts = 0, -1
            audio_timebase = Fraction(0, 1)
332
            video_timebase = Fraction(info.video_timebase.numerator, info.video_timebase.denominator)
333
            if info.has_audio:
334
335
336
                audio_timebase = Fraction(info.audio_timebase.numerator, info.audio_timebase.denominator)
                audio_start_pts = pts_convert(video_start_pts, video_timebase, audio_timebase, math.floor)
                audio_end_pts = pts_convert(video_end_pts, video_timebase, audio_timebase, math.ceil)
337
                audio_fps = info.audio_sample_rate
338
339
            video, audio, info = _read_video_from_file(
                video_path,
340
341
342
                video_width=self._video_width,
                video_height=self._video_height,
                video_min_dimension=self._video_min_dimension,
343
                video_max_dimension=self._video_max_dimension,
344
                video_pts_range=(video_start_pts, video_end_pts),
345
                video_timebase=video_timebase,
346
                audio_samples=self._audio_samples,
347
                audio_channels=self._audio_channels,
348
349
350
                audio_pts_range=(audio_start_pts, audio_end_pts),
                audio_timebase=audio_timebase,
            )
351
352
353
354
355

            info = {"video_fps": video_fps}
            if audio_fps is not None:
                info["audio_fps"] = audio_fps

356
357
358
359
360
361
        if self.frame_rate is not None:
            resampling_idx = self.resampling_idxs[video_idx][clip_idx]
            if isinstance(resampling_idx, torch.Tensor):
                resampling_idx = resampling_idx - resampling_idx[0]
            video = video[resampling_idx]
            info["video_fps"] = self.frame_rate
362
        assert len(video) == self.num_frames, f"{video.shape} x {self.num_frames}"
363
        return video, audio, info, video_idx
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407

    def __getstate__(self):
        video_pts_sizes = [len(v) for v in self.video_pts]
        # To be back-compatible, we convert data to dtype torch.long as needed
        # because for empty list, in legacy implementation, torch.as_tensor will
        # use torch.float as default dtype. This happens when decoding fails and
        # no pts is returned in the list.
        video_pts = [x.to(torch.int64) for x in self.video_pts]
        # video_pts can be an empty list if no frames have been decoded
        if video_pts:
            video_pts = torch.cat(video_pts)
            # avoid bug in https://github.com/pytorch/pytorch/issues/32351
            # TODO: Revert it once the bug is fixed.
            video_pts = video_pts.numpy()

        # make a copy of the fields of self
        d = self.__dict__.copy()
        d["video_pts_sizes"] = video_pts_sizes
        d["video_pts"] = video_pts
        # delete the following attributes to reduce the size of dictionary. They
        # will be re-computed in "__setstate__()"
        del d["clips"]
        del d["resampling_idxs"]
        del d["cumulative_sizes"]

        # for backwards-compatibility
        d["_version"] = 2
        return d

    def __setstate__(self, d):
        # for backwards-compatibility
        if "_version" not in d:
            self.__dict__ = d
            return

        video_pts = torch.as_tensor(d["video_pts"], dtype=torch.int64)
        video_pts = torch.split(video_pts, d["video_pts_sizes"], dim=0)
        # don't need this info anymore
        del d["video_pts_sizes"]

        d["video_pts"] = video_pts
        self.__dict__ = d
        # recompute attributes "clips", "resampling_idxs" and other derivative ones
        self.compute_clips(self.num_frames, self.step, self.frame_rate)