51 lines
1.9 KiB
Python
51 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 for the videos TODO: Intrinsics per video
|
|
:param allowed_extensions: Allowed video extensions, to not accidentally pick files that aren't videos
|
|
:return: generator that yields dict of {frames: [frame1, frame2, frame3], intrinsics: [fx, fy, tx, ty]}
|
|
"""
|
|
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, intrinsics):
|
|
# Single video file
|
|
video = cv2.VideoCapture(video_file)
|
|
|
|
try:
|
|
# Buffer to store 3 frames, yield when this fills up
|
|
current_frames = []
|
|
while video.grab():
|
|
current_frames.append(video.retrieve())
|
|
if len(current_frames) == 3:
|
|
temp_frames = current_frames
|
|
current_frames = []
|
|
# TODO: Consider converting frames to tensor
|
|
yield {'frames': temp_frames, 'intrinsics': intrinsics}
|
|
finally:
|
|
video.release()
|
|
|
|
|
|
def image_generator(root_folder):
|
|
"""
|
|
Create an image generator for unsupervised training
|
|
:param root_folder:
|
|
:return:
|
|
"""
|
|
pass
|