dataset

1 files changed, 14 insertions, 1 deletions

diff --git a/run.py b/run.py
index 00535fa..c58afa2 100644
--- a/run.py
+++ b/run.py
@@ -28,6 +28,7 @@ arguments_aOffset = 0
 arguments_bOffset = 0
 arguments_mixVideos = False
 arguments_averageVideos = False
+arguments_mixImages = False
 arguments_endpoint = None
 arguments_dataset = "dataset"
 
@@ -53,6 +54,8 @@ for strOption, strArgument in getopt.getopt(sys.argv[1:], '', [ strParameter[2:]
 		arguments_smooth = bool(strArgument)
 	elif strOption == '--mix-videos':
 		arguments_mixVideos = bool(strArgument)
+	elif strOption == '--mix-images':
+		arguments_mixImages = bool(strArgument)
 	elif strOption == '--average-videos':
 		arguments_averageVideos = bool(strArgument)
 	elif strOption == '--a-offset':
@@ -216,7 +219,6 @@ def store_frames(frames, outputPath, inputFirst=None, inputSecond=None, endpoint
 	    endpoint
 	  ])
 
-
 def tensor_to_image(np_val):
 	return (numpy.rollaxis(np_val, 0, 3)[:,:,::-1] * 255.0).astype(numpy.uint8)
 
@@ -264,6 +266,17 @@ 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(a_offset))))
+	inputSecond = load_image(os.path.join(arguments_strSecond, "frame_{:0>5}.png".format(int(b_offset))))
+	outputPath = './renders/' + arguments_strVideoOut
+	frames = process_tree(moduleNetwork, inputFirst, inputSecond, tensorOutput, arguments_steps * arguments_dilate, arguments_dilate)
+	frames = smooth_frames(moduleNetwork, tensorOutput, frames, arguments_smooth)
+	print("dilate... {}".format(arguments_dilate))
+	frames = dilate_frames(moduleNetwork, tensorOutput, frames, arguments_dilate)
+	store_frames(frames, outputPath, inputFirst, inputSecond, endpoint=arguments_endpoint, dataset=arguments_dataset)
+
 else:
 	# Morph two images
 	print("{} => {}".format(arguments_strFirst, arguments_strSecond))