53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
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
|