Update pintar.py

pintar.py

@@ -67,43 +67,35 @@ if __name__ == "__main__":
     colorUNet.load_state_dict(ckpt["colorUNet"])
     colorUNet.eval()
 
-    input_folder = args.input_folder
-    output_folder = args.output_folder
-    reference_image_path = args.reference_image
+    reference_img = Image.open(args.reference_image).convert("RGB")
+    reference_img, reference_img_lab = preprocessing(reference_img)
+    reference_img = reference_img.to(device)
+    reference_img_lab = reference_img_lab.to(device)
 
-    for root, dirs, files in os.walk(input_folder):
+    for root, dirs, files in os.walk(args.input_folder):
         for file in files:
             if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                 input_image_path = os.path.join(root, file)
 
-                img1 = Image.open(reference_image_path).convert("RGB")
-                width, height = img1.size
-                img2 = Image.open(input_image_path).convert("RGB")
-
-                img1, img1_lab = preprocessing(img1)
-                img2, img2_lab = preprocessing(img2)
-
-                img1 = img1.to(device)
-                img1_lab = img1_lab.to(device)
-                img2 = img2.to(device)
-                img2_lab = img2_lab.to(device)
+                img, img_lab = preprocessing(Image.open(input_image_path).convert("RGB"))
+                img = img.to(device)
+                img_lab = img_lab.to(device)
 
                 with torch.no_grad():
-                    img2_resize = F.interpolate(img2 / 255., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
-                    img1_L_resize = F.interpolate(img1_lab[:, :1, :, :] / 50., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
+                    img_resize = F.interpolate(img / 255., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
+                    img_L_resize = F.interpolate(img_lab[:, :1, :, :] / 50., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
 
-                    color_vector = colorEncoder(img2_resize)
-                    fake_ab = colorUNet((img1_L_resize, color_vector))
-                    fake_ab = F.interpolate(fake_ab * 110, size=(height, width), mode='bilinear', recompute_scale_factor=False, align_corners=False)
+                    color_vector = colorEncoder(img_resize)
+                    fake_ab = colorUNet((img_L_resize, color_vector))
+                    fake_ab = F.interpolate(fake_ab * 110, size=(img.size(2), img.size(3)), mode='bilinear', recompute_scale_factor=False, align_corners=False)
 
-                    fake_img = torch.cat((img1_lab[:, :1, :, :], fake_ab), 1)
+                    fake_img = torch.cat((img_lab[:, :1, :, :], fake_ab), 1)
                     fake_img = Lab2RGB_out(fake_img)
 
-                    relative_path = os.path.relpath(input_image_path, input_folder)
-                    output_subfolder = os.path.join(output_folder, os.path.dirname(relative_path), 'color')
+                    relative_path = os.path.relpath(input_image_path, args.input_folder)
+                    output_subfolder = os.path.join(args.output_folder, os.path.dirname(relative_path), 'color')
                     mkdirs(output_subfolder)
                     output_image_path = os.path.join(output_subfolder, f'{os.path.splitext(os.path.basename(input_image_path))[0]}_colorized.png')
                     io.imsave(output_image_path, fake_img)
 
-    print(f'Colored images have been saved to: {output_folder}')
-
+    print(f'Colored images have been saved to: {args.output_folder}')