diff --git a/unsupervised/load.py b/unsupervised/load.py new file mode 100644 index 0000000..f773f6f --- /dev/null +++ b/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