add adjustable size, random sampling from style ds

test_api.py

@@ -39,10 +39,11 @@ 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_dataset_pth, out_dir, alpha=1.0, vgg_pth='vgg_normalized.pth', decoder_pth='decoder.pth'):
+def run_adain(content_dir, style_dataset_pth, out_dir, alpha=1.0, dataset_size=100, vgg_pth='vgg_normalized.pth', decoder_pth='decoder.pth'):
 	content_pths = [Path(f) for f in glob(content_dir+'/*')]
+	num_content_imgs = len(content_pths)
 
-	assert len(content_pths) > 0, 'Failed to load content image'
+	assert num_content_imgs > 0, 'Failed to load content image'
 
 	# Load AdaIN model
 	vgg = torch.load(vgg_pth)
@@ -57,17 +58,22 @@ 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
 
-	for style_idx, style_item in enumerate(style_ds):
-		style_img = style_item['image']
-		if style_img.mode not in ("RGB", "L"):
-			style_img = style_img.convert("RGB")
-		style_tensor = t(style_img).unsqueeze(0).to(device)
+	if num_content_imgs * len(style_ds) > dataset_size:
+		num_style_per_content = int(np.ceil(dataset_size / num_content_imgs))
+	else:
+		num_style_per_content = len(style_ds)
 
-		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)
+		indices = random.sample(range(len(style_ds)), num_style_per_content)
+
+		for idx in indices:
+			style_img = style_ds[idx]['image']
+			if style_img.mode not in ("RGB", "L"):
+				style_img = style_img.convert("RGB")
+			style_tensor = t(style_img).unsqueeze(0).to(device)
     
 			# Start time
 			tic = time.perf_counter()
@@ -83,7 +89,7 @@ def run_adain(content_dir, style_dataset_pth, out_dir, alpha=1.0, vgg_pth='vgg_n
 			times.append(toc-tic)
             
 			# Save image
-			out_pth = out_dir + content_pth.stem + '_style_' + str(style_idx) + '_alpha' + str(alpha)
+			out_pth = out_dir + content_pth.stem + '_style_' + str(idx) + '_alpha' + str(alpha)
 			out_pth += content_pth.suffix
 			save_image(out_tensor, out_pth)