Start adding generators for unsupervised training

unsupervised/load.py

@@ -0,0 +1,52 @@
+import os
+
+import cv2
+
+
+def video_generator(video_path_or_folder, intrinsics, allowed_extensions=('mp4', 'mkv', 'mov')):
+    """
+    Create a generator for unsupervised training on depth sequences from a video file or folder of video files
+    :param video_path_or_folder: Video file or folder with list of video files to iterate through
+    :param intrinsics: Intrinsics should be static for a single video
+    :param allowed_extensions: Allowed video extensions, to not accidentally pick files that aren't videos
+    :return:
+    """
+    if os.path.isfile(video_path_or_folder):
+        # TODO: How to re-yield? Is this enough, since I'm just returning the actual generator?
+        #  Or do I need to iterate like below?
+        return _single_video_generator(video_path_or_folder)
+    else:
+        for root, dirs, files in os.walk(video_path_or_folder):
+            for file in files:
+                if os.path.splitext(file)[1] in allowed_extensions:
+                    for frames in _single_video_generator(file):
+                        yield frames
+
+
+def _single_video_generator(video_file):
+    # Single video file
+    video = cv2.VideoCapture(video_file)
+
+    try:
+        # Buffer to store 3 frames, yield when this fills up
+        current_frames = []
+        while video.grab():
+            # TODO: Should I be skipping frames or doing some magic to the frames? Or just leave that to some other
+            #  function assuming this will be used in a tf.data object?
+            current_frames.append(video.retrieve())
+            if len(current_frames) == 3:
+                temp_frames = current_frames
+                current_frames = []
+                # TODO: Convert to tensor or something notable (e.g. dict of 3 frames)
+                yield temp_frames
+    finally:
+        video.release()
+
+
+def image_generator(root_folder):
+    """
+    Create an image generator for unsupervised training
+    :param root_folder:
+    :return:
+    """
+    pass