Update style_transfer_folder.py

style_transfer_folder.py

@@ -58,35 +58,29 @@ if __name__ == '__main__':
     psp_encoder = PSPEncoder(args.psp_encoder_ckpt, output_size=args.size).to(device)
     psp_encoder.eval()
 
-    input_img_paths = sorted(glob.glob(os.path.join(args.input_img_path, '*.*')))
-    style_img_paths = sorted(glob.glob(os.path.join(args.style_img_path, '*.*')))[:]
-
     num = 0
 
-    for input_img_path in input_img_paths:
-        print(num)
-        num += 1
-
-        name_in = os.path.splitext(os.path.basename(input_img_path))[0]
-        img_in = cv2.imread(input_img_path, 1)
-        img_in_ten = cv2ten(img_in, device)
-        img_in = cv2.resize(img_in, (args.size, args.size))
-
-        for style_img_path in style_img_paths:
-            name_style = os.path.splitext(os.path.basename(style_img_path))[0]
-            img_style = cv2.imread(style_img_path, 1)
-            img_style_ten = cv2ten(img_style, device)
-            img_style = cv2.resize(img_style, (args.size, args.size))
-
-            with torch.no_grad():
-                sample_style = g_ema.get_z_embed(img_style_ten)
-                sample_in = psp_encoder(img_in_ten)
-                img_out_ten, _ = g_ema([sample_in], z_embed=sample_style, add_weight_index=args.add_weight_index,
-                                       input_is_latent=True, return_latents=False, randomize_noise=False)
-                img_out = ten2cv(img_out_ten)
-            out = np.concatenate([img_in, img_style, img_out], axis=1)
-            # out = img_out
-            cv2.imwrite(f'{args.outdir}/{name_in}_v_{name_style}.jpg', out)
+   
+    print(num)
+    num += 1
+
+    img_in = cv2.imread(args.input_img_path)
+    img_in_ten = cv2ten(img_in, device)
+    img_in = cv2.resize(img_in, (args.size, args.size))
+
+  
+    img_style = cv2.imread(args.style_img_path)
+    img_style_ten = cv2ten(img_style, device)
+    img_style = cv2.resize(img_style, (args.size, args.size))
+
+    with torch.no_grad():
+        sample_style = g_ema.get_z_embed(img_style_ten)
+        sample_in = psp_encoder(img_in_ten)
+        img_out_ten, _ = g_ema([sample_in], z_embed=sample_style, add_weight_index=args.add_weight_index,
+                               input_is_latent=True, return_latents=False, randomize_noise=False)
+        img_out = ten2cv(img_out_ten)
+    out = np.concatenate([img_in, img_style, img_out], axis=1)
+    cv2.imwrite('out.jpg', out)
 
     print('Done!')