vato007/fast-depth-tf
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix up generator to include intrinsics

Browse Source
This commit is contained in:
Piv
2021-08-29 19:26:15 +09:30
parent 90b73bf420
commit 2bb37b2722
1 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
unsupervised/load.py
View File

@@ -7,9 +7,9 @@ def video_generator(video_path_or_folder, intrinsics, allowed_extensions=('mp4',
""" """
Create a generator for unsupervised training on depth sequences from a video file or folder of video files 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 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 intrinsics: Intrinsics for the videos TODO: Intrinsics per video
:param allowed_extensions: Allowed video extensions, to not accidentally pick files that aren't videos :param allowed_extensions: Allowed video extensions, to not accidentally pick files that aren't videos
:return: :return: generator that yields dict of {frames: [frame1, frame2, frame3], intrinsics: [fx, fy, tx, ty]}
""" """
if os.path.isfile(video_path_or_folder): if os.path.isfile(video_path_or_folder):
# TODO: How to re-yield? Is this enough, since I'm just returning the actual generator? # TODO: How to re-yield? Is this enough, since I'm just returning the actual generator?
@@ -23,7 +23,7 @@ def video_generator(video_path_or_folder, intrinsics, allowed_extensions=('mp4',
yield frames yield frames
def _single_video_generator(video_file): def _single_video_generator(video_file, intrinsics):
# Single video file # Single video file
video = cv2.VideoCapture(video_file) video = cv2.VideoCapture(video_file)
@@ -31,14 +31,12 @@ def _single_video_generator(video_file):
# Buffer to store 3 frames, yield when this fills up # Buffer to store 3 frames, yield when this fills up
current_frames = [] current_frames = []
while video.grab(): 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()) current_frames.append(video.retrieve())
if len(current_frames) == 3: if len(current_frames) == 3:
temp_frames = current_frames temp_frames = current_frames
current_frames = [] current_frames = []
# TODO: Convert to tensor or something notable (e.g. dict of 3 frames) # TODO: Consider converting frames to tensor
yield temp_frames yield {'frames': temp_frames, 'intrinsics': intrinsics}
finally: finally:
video.release() video.release()