ucf101.py 5.08 KB
Newer Older
1
import os
2
from typing import Any, Dict, List, Tuple, Optional, Callable
3

4
from torch import Tensor
5

6
from .folder import find_classes, make_dataset
7
from .video_utils import VideoClips
8
9
10
11
from .vision import VisionDataset


class UCF101(VisionDataset):
12
    """
13
    `UCF101 <https://www.crcv.ucf.edu/data/UCF101.php>`_ dataset.
14
15
16
17

    UCF101 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
Bruno Korbar's avatar
Bruno Korbar committed
18
19
20
    ``step_between_clips``. The dataset itself can be downloaded from the dataset website;
    annotations that ``annotation_path`` should be pointing to can be downloaded from `here
    <https://www.crcv.ucf.edu/data/UCF101/UCF101TrainTestSplits-RecognitionTask.zip>`.
21
22
23
24
25
26
27
28
29
30
31

    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:
        root (string): Root directory of the UCF101 Dataset.
Bruno Korbar's avatar
Bruno Korbar committed
32
33
        annotation_path (str): path to the folder containing the split files;
            see docstring above for download instructions of these files
34
35
36
37
38
39
40
41
42
        frames_per_clip (int): number of frames in a clip.
        step_between_clips (int, optional): number of frames between each clip.
        fold (int, optional): which fold to use. Should be between 1 and 3.
        train (bool, optional): if ``True``, creates a dataset from the train split,
            otherwise from the ``test`` split.
        transform (callable, optional): A function/transform that  takes in a TxHxWxC video
            and returns a transformed version.

    Returns:
43
44
45
46
47
48
        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
49
    """
50

51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
    def __init__(
        self,
        root: str,
        annotation_path: str,
        frames_per_clip: int,
        step_between_clips: int = 1,
        frame_rate: Optional[int] = None,
        fold: int = 1,
        train: bool = True,
        transform: Optional[Callable] = None,
        _precomputed_metadata: Optional[Dict[str, Any]] = None,
        num_workers: int = 1,
        _video_width: int = 0,
        _video_height: int = 0,
        _video_min_dimension: int = 0,
66
        _audio_samples: int = 0,
67
    ) -> None:
68
        super().__init__(root)
69
        if not 1 <= fold <= 3:
70
            raise ValueError(f"fold should be between 1 and 3, got {fold}")
71

72
        extensions = ("avi",)
73
74
75
        self.fold = fold
        self.train = train

76
        self.classes, class_to_idx = find_classes(self.root)
77
78
        self.samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file=None)
        video_list = [x[0] for x in self.samples]
79
80
81
82
83
84
        video_clips = VideoClips(
            video_list,
            frames_per_clip,
            step_between_clips,
            frame_rate,
            _precomputed_metadata,
85
86
87
88
89
            num_workers=num_workers,
            _video_width=_video_width,
            _video_height=_video_height,
            _video_min_dimension=_video_min_dimension,
            _audio_samples=_audio_samples,
90
        )
91
92
93
94
        # we bookkeep the full version of video clips because we want to be able
        # to return the meta data of full version rather than the subset version of
        # video clips
        self.full_video_clips = video_clips
95
96
        self.indices = self._select_fold(video_list, annotation_path, fold, train)
        self.video_clips = video_clips.subset(self.indices)
97
        self.transform = transform
98

99
    @property
100
    def metadata(self) -> Dict[str, Any]:
101
        return self.full_video_clips.metadata
102

103
    def _select_fold(self, video_list: List[str], annotation_path: str, fold: int, train: bool) -> List[int]:
104
        name = "train" if train else "test"
105
        name = f"{name}list{fold:02d}.txt"
106
        f = os.path.join(annotation_path, name)
107
        selected_files = set()
108
        with open(f) as fid:
109
            data = fid.readlines()
110
111
112
            data = [x.strip().split(" ")[0] for x in data]
            data = [os.path.join(self.root, x) for x in data]
            selected_files.update(data)
113
        indices = [i for i in range(len(video_list)) if video_list[i] in selected_files]
114
115
        return indices

116
    def __len__(self) -> int:
117
118
        return self.video_clips.num_clips()

119
    def __getitem__(self, idx: int) -> Tuple[Tensor, Tensor, int]:
120
        video, audio, info, video_idx = self.video_clips.get_clip(idx)
121
        label = self.samples[self.indices[video_idx]][1]
122

123
124
125
        if self.transform is not None:
            video = self.transform(video)

126
        return video, audio, label