padding

1 files changed, 8 insertions, 4 deletions

diff --git a/run.py b/run.py
index 1568a9d..b059835 100644
--- a/run.py
+++ b/run.py
@@ -84,15 +84,17 @@ def recurse_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, count,
   a_fn = os.path.join(a, "frame_{:0>5}.png".format(int(frame_index + a_offset)))
   b_fn = os.path.join(b, "frame_{:0>5}.png".format(int(frame_index + b_offset)))
   print("{} => {}".format(a_fn, b_fn))
+  if not os.path.exists(a_fn) or not os.path.exists(b_fn):
+  	return [ None ]
   a_np = load_image(a_fn)
   b_np = load_image(b_fn)
   img_np = recurse_two_frames(moduleNetwork, tensorOutput, a_np, b_np, frame_index, count / 2, count / 4)
   if step < 2 * dilate:
-    return [img_np]
+    return [ img_np ]
   else:
     left = recurse_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, count, step, frame_index - (step/2), dilate)
     right = recurse_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, count, step, frame_index + (step/2), dilate)
-    return left + [img_np] + right
+    return left + [ img_np ] + right
 
 def process_two_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, steps, dilate):
   steps *= dilate
@@ -109,6 +111,8 @@ def average_two_videos(moduleNetwork, tensorOutput, a, b, a_offset, b_offset, st
     a_fn = os.path.join(a, "frame_{:0>5}.png".format(int(i + a_offset)))
     b_fn = os.path.join(b, "frame_{:0>5}.png".format(int(i + b_offset)))
     print("{} => {}".format(a_fn, b_fn))
+    if not os.path.exists(a_fn) and not os.path.exists(b_fn):
+    	continue
     a_np = load_image(a_fn)
     b_np = load_image(b_fn)
     tensorInputFirst = torch.FloatTensor(a_np)
@@ -207,8 +211,8 @@ def pad_frames(writer, base_path, start, end):
   for index in range(start, end):
     fn = os.path.join(base_path, "frame_{:0>5}.png".format(int(index + start)))
     if os.path.exists(fn):
-      PIL.Image.open(fn)
-      writer.write_frame(np)
+      img = PIL.Image.open(fn)
+      writer.write_frame(img)
 
 def tensor_to_image(np_val):
   return (numpy.rollaxis(np_val, 0, 3)[:,:,::-1] * 255.0).astype(numpy.uint8)