Update test_api.py

test_api.py

@@ -40,10 +40,10 @@ def style_transfer(content_tensor, style_tensor, encoder, decoder, alpha=1.0):
 	return decoder(mix_enc)
 
 def run_adain(content_dir, style_dataset_pth, out_dir, alpha=1.0, vgg_pth='vgg_normalized.pth', decoder_pth='decoder.pth'):
-    content_pths = [Path(f) for f in glob(content_dir+'/*')]
-    style_pths = [Path(f) for f in glob(style_dir+'/*')]
+	content_pths = [Path(f) for f in glob(content_dir+'/*')]
+	style_pths = [Path(f) for f in glob(style_dir+'/*')]
 
-    assert len(content_pths) > 0, 'Failed to load content image'
+	assert len(content_pths) > 0, 'Failed to load content image'
 	assert len(style_pths) > 0, 'Failed to load style image'
 
 	# Load AdaIN model
@@ -59,34 +59,34 @@ def run_adain(content_dir, style_dataset_pth, out_dir, alpha=1.0, vgg_pth='vgg_n
 	times = []
 
     style_ds = load_dataset(style_dataset_pth, split="train")
-    # do i need to stick a dataloader around this? idk
+	# do i need to stick a dataloader around this? idk
 
     for style_idx, style_item in enumerate(style_ds):
-        style_img = style_item['image']
-        print(style_img)
-        style_tensor = t(style_img).unsqueeze(0).to(device)
+		style_img = style_item['image']
+		print(style_img)
+		style_tensor = t(style_img).unsqueeze(0).to(device)
 
-    	for content_pth in content_pths:
-    		content_img = Image.open(content_pth)
-    		content_tensor = t(content_img).unsqueeze(0).to(device)
+		for content_pth in content_pths:
+			content_img = Image.open(content_pth)
+			content_tensor = t(content_img).unsqueeze(0).to(device)
     
-            # Start time
-            tic = time.perf_counter()
+			# Start time
+			tic = time.perf_counter()
             
-            # Execute style transfer
-            with torch.no_grad():
-                out_tensor = style_transfer(content_tensor, style_tensor, model.encoder, model.decoder, alpha).cpu()
+			# Execute style transfer
+			with torch.no_grad():
+				out_tensor = style_transfer(content_tensor, style_tensor, model.encoder, model.decoder, alpha).cpu()
         
-             # End time
-            toc = time.perf_counter()
-            print("Content: " + content_pth.stem + ". Style: " \
-                + str(style_idx) + '. Alpha: ' + str(alpha) + '. Style Transfer time: %.4f seconds' % (toc-tic))
-            times.append(toc-tic)
+			# End time
+			toc = time.perf_counter()
+			print("Content: " + content_pth.stem + ". Style: " \
+				+ str(style_idx) + '. Alpha: ' + str(alpha) + '. Style Transfer time: %.4f seconds' % (toc-tic))
+			times.append(toc-tic)
             
-            # Save image
-            out_pth = out_dir + content_pth.stem + '_style_' + str(style_idx) + '_alpha' + str(alpha)
-            out_pth += content_pth.suffix
-            save_image(out_tensor, out_pth)
+			# Save image
+			out_pth = out_dir + content_pth.stem + '_style_' + str(style_idx) + '_alpha' + str(alpha)
+			out_pth += content_pth.suffix
+			save_image(out_tensor, out_pth)
 
 	# Remove runtime of first iteration because it is flawed for some unknown reason
 	if len(times) > 1: