From 2bb37b2722769a9827eef9a4358bcad87340ffaf Mon Sep 17 00:00:00 2001
From: Piv <18462828+Piv200@users.noreply.github.com>
Date: Sun, 29 Aug 2021 19:26:15 +0930
Subject: [PATCH] Fix up generator to include intrinsics

---
 unsupervised/load.py | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/unsupervised/load.py b/unsupervised/load.py
index f773f6f..48c9c69 100644
--- a/unsupervised/load.py
+++ b/unsupervised/load.py
@@ -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
     :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
-    :return:
+    :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?
@@ -23,7 +23,7 @@ def video_generator(video_path_or_folder, intrinsics, allowed_extensions=('mp4',
                         yield frames
 
 
-def _single_video_generator(video_file):
+def _single_video_generator(video_file, intrinsics):
     # Single 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
         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
+                # TODO: Consider converting frames to tensor
+                yield {'frames': temp_frames, 'intrinsics': intrinsics}
     finally:
         video.release()