updated to check for existence of style transferred image

test.py

@@ -19,6 +19,7 @@ parser.add_argument('--style_dir', type=str, help='Content image folder path')
 parser.add_argument('--decoder_weight', type=str, default='decoder.pth', help='Decoder weight file path')
 parser.add_argument('--alpha', type=float, default=1.0, choices=[Range(0.0, 1.0)], help='Alpha [0.0, 1.0] controls style transfer level')
 parser.add_argument('--cuda', action='store_true', help='Use CUDA')
+parser.add_argument('--output_dir', type=str, default="results")
 parser.add_argument('--grid_pth', type=str, default=None, help='Specify a grid image path (default=None) if generate a grid image that contains all style transferred images')
 parser.add_argument('--color_control', action='store_true', help='Preserve content color')
 args = parser.parse_args()
@@ -71,13 +72,12 @@ def main():
 	assert len(style_pths) > 0, 'Failed to load style image'
 
 	# Prepare directory for saving results
-	out_dir = './results/'
-	os.makedirs(out_dir, exist_ok=True)
+	os.makedirs(args.output_dir, exist_ok=True)
 
 	# Load AdaIN model
-	vgg = torch.load('vgg_normalized.pth')
+	vgg = torch.load('vgg_normalized.pth', weights_only=False)
 	model = AdaINNet(vgg).to(device)
-	model.decoder.load_state_dict(torch.load(args.decoder_weight))
+	model.decoder.load_state_dict(torch.load(args.decoder_weight, weights_only=False))
 	model.eval()
 	
 	# Prepare image transform
@@ -95,14 +95,27 @@ def main():
 
 	for content_pth in content_pths:
 		content_img = Image.open(content_pth)
+		if not content_img.mode == "RGB":
+			content_img = content_img.convert("RGB")
 		content_tensor = t(content_img).unsqueeze(0).to(device)
 		
 		if args.grid_pth:
 			imgs.append(content_img)
 
 		for style_pth in style_pths:
+
+			# check if style transferred image exists already
+			out_pth = os.path.join(args.output_dir, content_pth.stem + '_style_' + style_pth.stem + '_alpha' + str(args.alpha) + content_pth.suffix)
+			if os.path.isfile(out_pth):
+				print("Skipping existing file")
+				continue
+
+			style_img = Image.open(style_pth)
+
+			if not style_img.mode == "RGB":
+				style_img = style_img.convert("RGB")
 			
-			style_tensor = t(Image.open(style_pth)).unsqueeze(0).to(device)
+			style_tensor = t(style_img).unsqueeze(0).to(device)
 			
 			# Linear Histogram Matching if needed
 			if args.color_control:
@@ -122,9 +135,6 @@ def main():
 			times.append(toc-tic)
 
 			# Save image
-			out_pth = out_dir + content_pth.stem + '_style_' + style_pth.stem + '_alpha' + str(args.alpha)
-			if args.color_control: out_pth += '_colorcontrol'
-			out_pth += content_pth.suffix
 			save_image(out_tensor, out_pth)
 
 			if args.grid_pth: