Update pintar.py

pintar.py

@@ -1,13 +1,12 @@
 import os
 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
+import argparse
 
 os.environ["CUDA_VISIBLE_DEVICES"] = '0'
 
@@ -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):
@@ -50,13 +49,19 @@ def preprocessing(inputs):
     return img.unsqueeze(0), img_lab.unsqueeze(0)
 
 if __name__ == "__main__":
-    device = "cuda"
+    parser = argparse.ArgumentParser(description="Colorize manga images.")
+    parser.add_argument("-i", "--input", type=str, required=True, help="Path to input image directory")
+    parser.add_argument("-r", "--reference", type=str, required=True, help="Path to reference image")
+    parser.add_argument("-o", "--output", type=str, required=True, help="Output directory")
+    parser.add_argument("-ckpt", "--checkpoint", type=str, required=True, help="Path to model checkpoint")
 
-    # Specify the paths here
-    img_path = 'path/to/your/input/image.jpg'
-    ckpt_path = 'path/to/your/model_checkpoint.pt'
-    reference_image_path = 'path/to/your/reference/image.jpg'
+    args = parser.parse_args()
 
+    device = "cuda"
+    input_image_dir = args.input
+    output_directory = args.output
+    model_checkpoint_path = args.checkpoint
+    reference_image_path = args.reference
     imgsize = 256
 
     ckpt = torch.load(ckpt_path, map_location=lambda storage, loc: storage)