Update test_api.py

test_api.py

@@ -10,10 +10,10 @@ from torchvision.utils import save_image
 from torchvision.transforms import ToPILImage
 from utils import adaptive_instance_normalization, grid_image, transform,linear_histogram_matching, Range
 from glob import glob
+from datasets import load_dataset
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-
 def style_transfer(content_tensor, style_tensor, encoder, decoder, alpha=1.0):
 	"""
 	Given content image and style image, generate feature maps with encoder, apply 
@@ -39,17 +39,13 @@ def style_transfer(content_tensor, style_tensor, encoder, decoder, alpha=1.0):
 	mix_enc = alpha * transfer_enc + (1-alpha) * content_enc
 	return decoder(mix_enc)
 
-def run_adain(content_dir, style_dir, vgg_pth='vgg_normalized.pth', decoder_pth='decoder.pth', alpha=1.0):
+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+'/*')]
 
     assert len(content_pths) > 0, 'Failed to load content image'
 	assert len(style_pths) > 0, 'Failed to load style image'
 
-	# Prepare directory for saving results
-	out_dir = tempfile.mkdtemp()
-	os.makedirs(out_dir, exist_ok=True)
-
 	# Load AdaIN model
 	vgg = torch.load(vgg_pth)
 	model = AdaINNet(vgg).to(device)
@@ -62,31 +58,35 @@ def run_adain(content_dir, style_dir, vgg_pth='vgg_normalized.pth', decoder_pth=
 	# Timer
 	times = []
 
-	for content_pth in content_pths:
-		content_img = Image.open(content_pth)
-		content_tensor = t(content_img).unsqueeze(0).to(device)
+    style_ds = load_dataset(style_dataset_pth, split="train")
+    # do i need to stick a dataloader around this? idk
 
-		for style_pth in style_pths:
-			
-			style_tensor = t(Image.open(style_pth)).unsqueeze(0).to(device)
+    for style_item in style_ds:
+        style_img = style_item['image']
+        print(style_img)
+        style_tensor = t(style_img).unsqueeze(0).to(device)
 
-			# 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()
-		
-			 # End time
-			toc = time.perf_counter()
-			print("Content: " + content_pth.stem + ". Style: " \
-				+ style_pth.stem + '. Alpha: ' + str(alpha) + '. Style Transfer time: %.4f seconds' % (toc-tic))
-			times.append(toc-tic)
+    	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()
+            
+            # 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: " \
+                + style_pth.stem + '. Alpha: ' + str(alpha) + '. Style Transfer time: %.4f seconds' % (toc-tic))
+            times.append(toc-tic)
             
-			# Save image
-			out_pth = out_dir + content_pth.stem + '_style_' + style_pth.stem + '_alpha' + str(alpha)
-			out_pth += content_pth.suffix
-			save_image(out_tensor, out_pth)
+            # Save image
+            out_pth = out_dir + content_pth.stem + '_style_' + style_pth.stem + '_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: