check for doubled frames and use correct sleep command

1 files changed, 7 insertions, 3 deletions

diff --git a/data/recursive_dataset.py b/data/recursive_dataset.py
index b5a4ed2..7bf7941 100644
--- a/data/recursive_dataset.py
+++ b/data/recursive_dataset.py
@@ -1,5 +1,4 @@
 import os
-import time
 import os.path
 from data.base_dataset import BaseDataset, get_transform
 from data.image_folder import make_dataset
@@ -20,8 +19,13 @@ class RecursiveDataset(BaseDataset):
     def __getitem__(self, index):
         A_path = os.path.join(self.opt.render_dir, "frame_{:05d}.png".format(index))
         # print(A_path)
-        while not os.path.exists(A_path):
-            time.sleep(self.opt.poll_delay)
+        if not os.path.exists(A_path):
+            A_path2 = os.path.join(self.opt.render_dir, "frame_{:05d}.png".format(index+1))
+            if os.path.exists(A_path2):
+                A_path = A_path2
+            else:
+                while not os.path.exists(A_path):
+                    gevent.sleep(self.opt.poll_delay)
         A_img = Image.open(A_path).convert('RGB')
         A = self.transform(A_img)
         if self.opt.which_direction == 'BtoA':