"examples/vscode:/vscode.git/clone" did not exist on "1fb1007716378822c97d902704b4b5a8015180be"
kinetics.py 3.5 KB
Newer Older
1
2
from .utils import list_dir
from .folder import make_dataset
3
from .video_utils import VideoClips
4
5
6
from .vision import VisionDataset


7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Kinetics400(VisionDataset):
    """
    `Kinetics-400 <https://deepmind.com/research/open-source/open-source-datasets/kinetics/>`_
    dataset.

    Kinetics-400 is an action recognition video dataset.
    This dataset consider every video as a collection of video clips of fixed size, specified
    by ``frames_per_clip``, where the step in frames between each clip is given by
    ``step_between_clips``.

    To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5``
    and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two
    elements will come from video 1, and the next three elements from video 2.
    Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all
    frames in a video might be present.

    Internally, it uses a VideoClips object to handle clip creation.

    Args:
26
27
28
29
30
31
32
33
34
35
36
37
        root (string): Root directory of the Kinetics-400 Dataset. Should be structured as follows:
            .. code::

                root/
                ├── class1
                │   ├── clip1.avi
                │   ├── clip2.avi
                │   └── ...
                └── class2
                    ├── clipx.avi
                    └── ...

38
39
40
41
42
43
        frames_per_clip (int): number of frames in a clip
        step_between_clips (int): number of frames between each clip
        transform (callable, optional): A function/transform that  takes in a TxHxWxC video
            and returns a transformed version.

    Returns:
44
45
46
47
48
49
        tuple: A 3-tuple with the following entries:

            - video (Tensor[T, H, W, C]): the `T` video frames
            - audio(Tensor[K, L]): the audio frames, where `K` is the number of channels
              and `L` is the number of points
            - label (int): class of the video clip
50
51
    """

52
    def __init__(self, root, frames_per_clip, step_between_clips=1, frame_rate=None,
53
54
                 extensions=('avi',), transform=None, _precomputed_metadata=None,
                 num_workers=1, _video_width=0, _video_height=0,
55
                 _video_min_dimension=0, _audio_samples=0, _audio_channels=0):
56
        super(Kinetics400, self).__init__(root)
57
58
59
60
61
62

        classes = list(sorted(list_dir(root)))
        class_to_idx = {classes[i]: i for i in range(len(classes))}
        self.samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file=None)
        self.classes = classes
        video_list = [x[0] for x in self.samples]
63
64
65
66
67
68
        self.video_clips = VideoClips(
            video_list,
            frames_per_clip,
            step_between_clips,
            frame_rate,
            _precomputed_metadata,
69
70
71
72
73
            num_workers=num_workers,
            _video_width=_video_width,
            _video_height=_video_height,
            _video_min_dimension=_video_min_dimension,
            _audio_samples=_audio_samples,
74
            _audio_channels=_audio_channels,
75
        )
76
        self.transform = transform
77

78
79
80
81
    @property
    def metadata(self):
        return self.video_clips.metadata

82
83
84
85
86
87
88
    def __len__(self):
        return self.video_clips.num_clips()

    def __getitem__(self, idx):
        video, audio, info, video_idx = self.video_clips.get_clip(idx)
        label = self.samples[video_idx][1]

89
90
91
        if self.transform is not None:
            video = self.transform(video)

92
        return video, audio, label