process sequences

1 files changed, 6 insertions, 4 deletions

diff --git a/run.py b/run.py
index 61f767f..60854d8 100644
--- a/run.py
+++ b/run.py
@@ -119,14 +119,16 @@ def process_tree(moduleNetwork, a, b, tensorOutput, steps):
 		right = process_tree(moduleNetwork, middle, b, tensorOutput, steps / 2)
 		return left + [np_middle] + right
 
-def store_frames(frames, outputPath):
+def store_frames(frames, outputPath, inputFirst=None, inputSecond=None):
 	print('writing {}'.format(outputPath))
 	print('frames: {}'.format(len(frames)))
 	writer = FFMPEG_VideoWriter(outputPath, (1024, 512), 25)
-	writer.write_frame(inputFirst)
+	if inputFirst is not None:
+		writer.write_frame(inputFirst)
 	for frame in frames:
 		writer.write_frame(frame)
-	writer.write_frame(inputSecond)
+	if inputSecond is not None:
+		writer.write_frame(inputSecond)
 
 def load_image(path):
 	return numpy.rollaxis(numpy.asarray(PIL.Image.open(path))[:,:,::-1], 2, 0).astype(numpy.float32) / 255.0
@@ -169,4 +171,4 @@ else:
 	inputSecond = load_image(arguments_strSecond)
 	frames = process_tree(moduleNetwork, inputFirst, inputSecond, tensorOutput, arguments_steps)
 	outputPath = './renders/' + arguments_strVideoOut
-	store_frames(frames, outputPath)
+	store_frames(frames, outputPath, inputFirst, inputSecond)