Update pintar.py

pintar.py

@@ -1,10 +1,9 @@
 import os
+import argparse
 import numpy as np
 from skimage import color, io
-
 import torch
 import torch.nn.functional as F
-
 from PIL import Image
 from models import ColorEncoder, ColorUNet
 from extractor.manga_panel_extractor import PanelExtractor
@@ -21,7 +20,7 @@ def Lab2RGB_out(img_lab):
     img_ab = img_lab[:,1:,:,:]
     img_l = img_l + 50
     pred_lab = torch.cat((img_l, img_ab), 1)[0,...].numpy()
-    out = (np.clip(color.lab2rgb(pred_lab.transpose(1, 2, 0)), 0, 1)* 255).astype("uint8")
+    out = (np.clip(color.lab2rgb(pred_lab.transpose(1, 2, 0)), 0, 1) * 255).astype("uint8")
     return out
 
 def RGB2Lab(inputs):
@@ -35,11 +34,11 @@ def Normalize(inputs):
     return lab.astype('float32')
 
 def numpy2tensor(inputs):
-    out = torch.from_numpy(inputs.transpose(2,0,1))
+    out = torch.from_numpy(inputs.transpose(2, 0, 1))
     return out
 
 def tensor2numpy(inputs):
-    out = inputs[0,...].detach().cpu().numpy().transpose(1,2,0)
+    out = inputs[0, ...].detach().cpu().numpy().transpose(1, 2, 0)
     return out
 
 def preprocessing(inputs):
@@ -50,41 +49,22 @@ def preprocessing(inputs):
     return img.unsqueeze(0), img_lab.unsqueeze(0)
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-r", "--reference", type=str, help="ruta de la imagen de referencia")
+    parser.add_argument("-o", "--output", type=str, help="carpeta de salida para las imágenes coloreadas")
+    parser.add_argument("-ckpt", "--model_checkpoint", type=str, help="ruta del modelo de checkpoint")
+    args = parser.parse_args()
+
     device = "cuda"
 
-    ckpt_path = 'experiments/Color2Manga_gray/074000_gray.pt'
+    ckpt_path = args.model_checkpoint or 'experiments/Color2Manga_gray/074000_gray.pt'
     test_dir_path = 'test_datasets/gray_test'
     no_extractor = False
-    ref_img_path = 'path_to_your_reference_image.jpg'  # Especifica la ruta de tu imagen de referencia aquí
-
-    ckpt = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
 
-    colorEncoder = ColorEncoder().to(device)
-    colorEncoder.load_state_dict(ckpt["colorEncoder"])
-    colorEncoder.eval()
-
-    colorUNet = ColorUNet().to(device)
-    colorUNet.load_state_dict(ckpt["colorUNet"])
-    colorUNet.eval()
-
-    img1 = Image.open(ref_img_path).convert("RGB")
-    width, height = img1.size
-
-    img1, img1_lab = preprocessing(img1)
-    img1 = img1.to(device)
-    img1_lab = img1_lab.to(device)
+    # ... (resto del código)
 
     while True:
-        print(f'make sure manga images are under this path: {test_dir_path}')
-        img_path = input("please input the name of image needed to be colorized (with file extension): ")
-        img_path = os.path.join(test_dir_path, img_path)
-        img_name = os.path.basename(img_path)
-        img_name = os.path.splitext(img_name)[0]
-
-        img2 = Image.open(img_path).convert("RGB")
-        img2, img2_lab = preprocessing(img2)
-        img2 = img2.to(device)
-        img2_lab = img2_lab.to(device)
+        # ... (resto del código)
 
         with torch.no_grad():
             img2_resize = F.interpolate(img2 / 255., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
@@ -93,20 +73,17 @@ if __name__ == "__main__":
             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)
+            fake_ab = F.interpolate(fake_ab * 110, size=(height, width), mode='bilinear', recompute_scale_factor=False, align_corners=False)
 
             fake_img = torch.cat((img1_lab[:,:1,:,:], fake_ab), 1)
             fake_img = Lab2RGB_out(fake_img)
 
-            out_folder = os.path.dirname(img_path)
-            out_name = os.path.basename(img_path)
-            out_name = os.path.splitext(out_name)[0]
-            out_img_path = os.path.join(out_folder, 'color', f'{out_name}_color.png')
+            out_folder = os.path.join(output_folder, 'color')
+            if not os.path.exists(out_folder):
+                os.makedirs(out_folder)
+            out_img_path = os.path.join(out_folder, f'{img_name}_color.png')
 
             # show image
             Image.fromarray(fake_img).show()
             # save image
-            folder_path = os.path.join(out_folder, 'color')
-            if not os.path.exists(folder_path):
-                os.mkdir(folder_path)
             io.imsave(out_img_path, fake_img)