Expose frame-rate and cache to video datasets (#1356)

pytorch · Sep 20, 2019 · 85ffd93 · 85ffd93
1 parent 31fad34
commit 85ffd93
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 6 deletions.
diff --git a/torchvision/datasets/hmdb51.py b/torchvision/datasets/hmdb51.py
@@ -50,7 +50,8 @@ class HMDB51(VisionDataset):
     }
 
     def __init__(self, root, annotation_path, frames_per_clip, step_between_clips=1,
-                 fold=1, train=True, transform=None):
+                 frame_rate=None, fold=1, train=True, transform=None,
+                 _precomputed_metadata=None):
         super(HMDB51, self).__init__(root)
         if not 1 <= fold <= 3:
             raise ValueError("fold should be between 1 and 3, got {}".format(fold))
@@ -64,7 +65,13 @@ def __init__(self, root, annotation_path, frames_per_clip, step_between_clips=1,
         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]
-        video_clips = VideoClips(video_list, frames_per_clip, step_between_clips)
+        video_clips = VideoClips(
+            video_list,
+            frames_per_clip,
+            step_between_clips,
+            frame_rate,
+            _precomputed_metadata,
+        )
         self.indices = self._select_fold(video_list, annotation_path, fold, train)
         self.video_clips = video_clips.subset(self.indices)
         self.transform = transform

diff --git a/torchvision/datasets/kinetics.py b/torchvision/datasets/kinetics.py
@@ -36,7 +36,8 @@ class Kinetics400(VisionDataset):
         label (int): class of the video clip
     """
 
-    def __init__(self, root, frames_per_clip, step_between_clips=1, transform=None):
+    def __init__(self, root, frames_per_clip, step_between_clips=1, frame_rate=None,
+                 extensions=('avi',), transform=None, _precomputed_metadata=None):
         super(Kinetics400, self).__init__(root)
         extensions = ('avi',)
 
@@ -45,7 +46,13 @@ def __init__(self, root, frames_per_clip, step_between_clips=1, transform=None):
         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]
-        self.video_clips = VideoClips(video_list, frames_per_clip, step_between_clips)
+        self.video_clips = VideoClips(
+            video_list,
+            frames_per_clip,
+            step_between_clips,
+            frame_rate,
+            _precomputed_metadata,
+        )
         self.transform = transform
 
     def __len__(self):

diff --git a/torchvision/datasets/ucf101.py b/torchvision/datasets/ucf101.py
@@ -43,7 +43,8 @@ class UCF101(VisionDataset):
     """
 
     def __init__(self, root, annotation_path, frames_per_clip, step_between_clips=1,
-                 fold=1, train=True, transform=None):
+                 frame_rate=None, fold=1, train=True, transform=None,
+                 _precomputed_metadata=None):
         super(UCF101, self).__init__(root)
         if not 1 <= fold <= 3:
             raise ValueError("fold should be between 1 and 3, got {}".format(fold))
@@ -57,7 +58,13 @@ def __init__(self, root, annotation_path, frames_per_clip, step_between_clips=1,
         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]
-        video_clips = VideoClips(video_list, frames_per_clip, step_between_clips)
+        video_clips = VideoClips(
+            video_list,
+            frames_per_clip,
+            step_between_clips,
+            frame_rate,
+            _precomputed_metadata,
+        )
         self.indices = self._select_fold(video_list, annotation_path, fold, train)
         self.video_clips = video_clips.subset(self.indices)
         self.transform = transform