Update test_api.py

test_api.py

@@ -61,7 +61,7 @@ def run_adain(content_dir, style_dataset_pth, out_dir, alpha=1.0, vgg_pth='vgg_n
     style_ds = load_dataset(style_dataset_pth, split="train")
     # do i need to stick a dataloader around this? idk
 
-    for style_item in style_ds:
+    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)
@@ -80,11 +80,11 @@ def run_adain(content_dir, style_dataset_pth, out_dir, alpha=1.0, vgg_pth='vgg_n
              # 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))
+                + 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_' + style_pth.stem + '_alpha' + str(alpha)
+            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)