dataset

1 files changed, 3 insertions, 2 deletions

diff --git a/run.py b/run.py
index 80b3ae9..850c313 100644
--- a/run.py
+++ b/run.py
@@ -199,11 +199,11 @@ def store_frames(frames, outputPath, inputFirst=None, inputSecond=None, endpoint
 	print('frames: {}'.format(len(frames)))
 	writer = FFMPEG_VideoWriter(outputPath, (1024, 512), 25)
 	if inputFirst is not None:
-		writer.write_frame(inputFirst)
+		writer.write_frame(tensor_to_image(inputFirst))
 	for frame in frames:
 		writer.write_frame(tensor_to_image(frame))
 	if inputSecond is not None:
-		writer.write_frame(inputSecond)
+		writer.write_frame(tensor_to_image(inputSecond))
 	writer.close()
 
 	if endpoint is not None:
@@ -266,6 +266,7 @@ elif arguments_steps == 0:
 	tensorInputSecond = load_image_tensor(arguments_strSecond)
 	process(moduleNetwork, tensorInputFirst, tensorInputSecond, tensorOutput)
 	PIL.Image.fromarray((numpy.rollaxis(tensorOutput.clamp(0.0, 1.0).numpy(), 0, 3)[:,:,::-1] * 255.0).astype(numpy.uint8)).save(arguments_strOut)
+
 elif arguments_mixImages:
 	print("{} => {}".format(arguments_strFirst, arguments_strSecond))
 	inputFirst = load_image(os.path.join(arguments_strFirst, "frame_{:0>5}.png".format(int(arguments_aOffset)+1)))