modules/pytorch_jacinto_ai/xvision/datasets/kinetics.py

   1 from .video_utils import VideoClips
   2 from .utils import list_dir
   3 from .folder import make_dataset
   4 from .vision import VisionDataset
   5
   6
   7 class Kinetics400(VisionDataset):
   8     """
   9     `Kinetics-400 <https://deepmind.com/research/open-source/open-source-datasets/kinetics/>`_
  10     dataset.
  11
  12     Kinetics-400 is an action recognition video dataset.
  13     This dataset consider every video as a collection of video clips of fixed size, specified
  14     by ``frames_per_clip``, where the step in frames between each clip is given by
  15     ``step_between_clips``.
  16
  17     To give an example, for 2 videos with 10 and 15 frames respectively, if ``frames_per_clip=5``
  18     and ``step_between_clips=5``, the dataset size will be (2 + 3) = 5, where the first two
  19     elements will come from video 1, and the next three elements from video 2.
  20     Note that we drop clips which do not have exactly ``frames_per_clip`` elements, so not all
  21     frames in a video might be present.
  22
  23     Internally, it uses a VideoClips object to handle clip creation.
  24
  25     Args:
  26         root (string): Root directory of the Kinetics-400 Dataset.
  27         frames_per_clip (int): number of frames in a clip
  28         step_between_clips (int): number of frames between each clip
  29         transform (callable, optional): A function/transform that  takes in a TxHxWxC video
  30             and returns a transformed version.
  31
  32     Returns:
  33         video (Tensor[T, H, W, C]): the `T` video frames
  34         audio(Tensor[K, L]): the audio frames, where `K` is the number of channels
  35             and `L` is the number of points
  36         label (int): class of the video clip
  37     """
  38
  39     def __init__(self, root, frames_per_clip, step_between_clips=1, transform=None):
  40         super(Kinetics400, self).__init__(root)
  41         extensions = ('avi',)
  42
  43         classes = list(sorted(list_dir(root)))
  44         class_to_idx = {classes[i]: i for i in range(len(classes))}
  45         self.samples = make_dataset(self.root, class_to_idx, extensions, is_valid_file=None)
  46         self.classes = classes
  47         video_list = [x[0] for x in self.samples]
  48         self.video_clips = VideoClips(video_list, frames_per_clip, step_between_clips)
  49         self.transform = transform
  50
  51     def __len__(self):
  52         return self.video_clips.num_clips()
  53
  54     def __getitem__(self, idx):
  55         video, audio, info, video_idx = self.video_clips.get_clip(idx)
  56         label = self.samples[video_idx][1]
  57
  58         if self.transform is not None:
  59             video = self.transform(video)
  60
  61         return video, audio, label