Add optional shuffling

Author: shi69

data/VideoFolder.py

@@ -2,9 +2,10 @@
 import torch
 import torch.utils.data as data
 
+from random import shuffle as list_shuffle  # for shuffling list
 from math import ceil
 from os import listdir
-from os.path import isdir, join
+from os.path import isdir, join, isfile
 from itertools import islice
 from numpy.core.multiarray import concatenate, ndarray
 from skvideo.io import FFmpegReader, ffprobe
@@ -68,7 +69,7 @@ def __call__(self, batch: iter) -> torch.Tensor or list(torch.Tensor):
 
 
 class VideoFolder(data.Dataset):
-    def __init__(self, root, transform=None, target_transform=None, video_index=False):
+    def __init__(self, root, transform=None, target_transform=None, video_index=False, shuffle=None):
         """
         Initialise a data.Dataset object for concurrent frame fetching from videos in a directory of folders of videos
 
@@ -79,12 +80,16 @@ def __init__(self, root, transform=None, target_transform=None, video_index=Fals
         :param target_transform: label transformation / mapping
         :type target_transform: object
         :param video_index: if True, the label will be the video index instead of target class
-        :type bool
+        :type video_index: bool
+        :param shuffle: None, 'init' or 'always'
+        :type shuffle: str
         """
         classes, class_to_idx = self._find_classes(root)
-        videos, frames, frames_per_video = self._make_data_set(root, classes, class_to_idx)
+        video_paths = self._find_videos(root, classes)
+        videos, frames, frames_per_video = self._make_data_set(root, video_paths, class_to_idx, shuffle == 'init')
 
         self.root = root
+        self.video_paths = video_paths
         self.videos = videos
         self.opened_videos = [[] for _ in videos]
         self.frames = frames
@@ -94,8 +99,14 @@ def __init__(self, root, transform=None, target_transform=None, video_index=Fals
         self.transform = transform
         self.target_transform = target_transform
         self.alternative_target = video_index
+        self.shuffle = shuffle
 
     def __getitem__(self, frame_idx):
+        if frame_idx == 0:
+            self.free()
+            if self.shuffle == 'always':
+                self._shuffle()
+
         frame_idx %= self.frames  # wrap around indexing, if asking too much
         video_idx = bisect(self.videos, ((frame_idx,),))  # video to which frame_idx belongs
         (last, first), (path, target) = self.videos[video_idx]  # get video metadata
@@ -138,11 +149,47 @@ def free(self):
         """
         Frees all video files' pointers
         """
+        print('Resetting data set internal state')
         for video in self.opened_videos:  # for every opened video
             for _ in range(len(video)):  # for as many times as pointers
                 opened_video = video.pop()  # pop an item
                 opened_video[2]._close()  # close the file
 
+    def _shuffle(self):
+        """
+        Shuffles the video list
+        by regenerating the sequence to sample sequentially
+        """
+        def _is_video_file(filename_):
+            return any(filename_.endswith(extension) for extension in VIDEO_EXTENSIONS)
+
+        root = self.root
+        video_paths = self.video_paths
+        class_to_idx = self.class_to_idx
+        list_shuffle(video_paths)  # shuffle
+
+        videos = list()
+        frames_per_video = list()
+        frames_counter = 0
+        for filename in tqdm(video_paths, ncols=80):
+            class_ = filename.split('/')[0]
+            data_path = join(root, filename)
+            if _is_video_file(data_path):
+                video_meta = ffprobe(data_path)
+                start_idx = frames_counter
+                frames = int(video_meta['video'].get('@nb_frames'))
+                frames_per_video.append(frames)
+                frames_counter += frames
+                item = ((frames_counter - 1, start_idx), (filename, class_to_idx[class_]))
+                videos.append(item)
+
+        sleep(0.5)  # allows for progress bar completion
+        # update the attributes with the altered sequence
+        self.video_paths = video_paths
+        self.videos = videos
+        self.frames = frames_counter
+        self.frames_per_video = frames_per_video
+
     @staticmethod
     def _find_classes(data_path):
         classes = [d for d in listdir(data_path) if isdir(join(data_path, d))]
@@ -151,25 +198,31 @@ def _find_classes(data_path):
         return classes, class_to_idx
 
     @staticmethod
-    def _make_data_set(data_path, classes, class_to_idx):
+    def _find_videos(root, classes):
+        return [join(c, d) for c in classes for d in listdir(join(root, c))]
+
+    @staticmethod
+    def _make_data_set(root, video_paths, class_to_idx, init_shuffle):
         def _is_video_file(filename_):
             return any(filename_.endswith(extension) for extension in VIDEO_EXTENSIONS)
 
+        if init_shuffle:
+            list_shuffle(video_paths)  # shuffle
+
         videos = list()
         frames_per_video = list()
         frames_counter = 0
-        for class_ in tqdm(classes, ncols=80):
-            class_path = join(data_path, class_)
-            for filename in listdir(class_path):
-                if _is_video_file(filename):
-                    video_path = join(class_path, filename)
-                    video_meta = ffprobe(video_path)
-                    start_idx = frames_counter
-                    frames = int(video_meta['video'].get('@nb_frames'))
-                    frames_per_video.append(frames)
-                    frames_counter += frames
-                    item = ((frames_counter - 1, start_idx), (join(class_, filename), class_to_idx[class_]))
-                    videos.append(item)
+        for filename in tqdm(video_paths, ncols=80):
+            class_ = filename.split('/')[0]
+            data_path = join(root, filename)
+            if _is_video_file(data_path):
+                video_meta = ffprobe(data_path)
+                start_idx = frames_counter
+                frames = int(video_meta['video'].get('@nb_frames'))
+                frames_per_video.append(frames)
+                frames_counter += frames
+                item = ((frames_counter - 1, start_idx), (filename, class_to_idx[class_]))
+                videos.append(item)
 
         sleep(0.5)  # allows for progress bar completion
         return videos, frames_counter, frames_per_video