Update pintar.py

pintar.py

@@ -1,12 +1,13 @@
 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'
 
@@ -20,7 +21,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):
@@ -49,16 +50,16 @@ def preprocessing(inputs):
     return img.unsqueeze(0), img_lab.unsqueeze(0)
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Colorize manga images.")
-    parser.add_argument("-i", "--input_folder", type=str, required=True, help="Path to the input folder containing manga images.")
-    parser.add_argument("-ckpt", "--model_checkpoint", type=str, required=True, help="Path to the model checkpoint file.")
-    parser.add_argument("-o", "--output_folder", type=str, required=True, help="Path to the output folder where colorized images will be saved.")
-    parser.add_argument("-ne", "--no_extractor", action="store_true", help="Do not segment the manga panels.")
-    args = parser.parse_args()
-
     device = "cuda"
 
-    ckpt = torch.load(args.model_checkpoint, map_location=lambda storage, loc: storage)
+    # 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'
+
+    imgsize = 256
+
+    ckpt = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
 
     colorEncoder = ColorEncoder().to(device)
     colorEncoder.load_state_dict(ckpt["colorEncoder"])
@@ -68,82 +69,32 @@ if __name__ == "__main__":
     colorUNet.load_state_dict(ckpt["colorUNet"])
     colorUNet.eval()
 
-    input_files = os.listdir(args.input_folder)
-
-    for input_file in input_files:
-        input_path = os.path.join(args.input_folder, input_file)
-
-        if os.path.isfile(input_path):
-            if args.no_extractor:
-                ref_img_path = input("Please enter the path of the reference image: ")
-
-                img1 = Image.open(ref_img_path).convert("RGB")
-                width, height = img1.size
-                img2 = Image.open(input_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)
-
-                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)
-
-                    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_img = torch.cat((img1_lab[:, :1, :, :], fake_ab), 1)
-                    fake_img = Lab2RGB_out(fake_img)
-
-                    out_folder = os.path.join(args.output_folder, 'color')
-                    mkdirs(out_folder)
-                    out_img_path = os.path.join(out_folder, f'{os.path.splitext(input_file)[0]}_color.png')
-                    io.imsave(out_img_path, fake_img)
-
-            else:
-                panel_extractor = PanelExtractor(min_pct_panel=5, max_pct_panel=90)  # You might need to adjust these parameters
-                panels, masks, panel_masks = panel_extractor.extract(input_path)
-                
-                ref_img_paths = []
-                print("Please enter the name of the reference image in order according to the number prompts on the picture")
-                for i in range(len(panels)):
-                    ref_img_path = input(f"{i+1}/{len(panels)} reference image:")
-                    ref_img_paths.append(ref_img_path)
-
-                fake_imgs = []
-                for i in range(len(panels)):
-                    img1 = Image.fromarray(panels[i]).convert("RGB")
-                    width, height = img1.size
-                    img2 = Image.open(ref_img_paths[i]).convert("RGB")
-
-                    img1, img1_lab = preprocessing(img1)
-                    img2, img2_lab = preprocessing(img2)
+    img_name = os.path.splitext(os.path.basename(img_path))[0]
+    img1 = Image.open(img_path).convert("RGB")
+    width, height = img1.size
+    img1, img1_lab = preprocessing(img1)
+    img2, img2_lab = preprocessing(Image.open(reference_image_path).convert("RGB"))
 
-                    img1 = img1.to(device)
-                    img1_lab = img1_lab.to(device)
-                    img2 = img2.to(device)
-                    img2_lab = img2_lab.to(device)
+    img1 = img1.to(device)
+    img1_lab = img1_lab.to(device)
+    img2 = img2.to(device)
+    img2_lab = img2_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)
+    with torch.no_grad():
+        img2_resize = F.interpolate(img2 / 255., size=(imgsize, imgsize), mode='bilinear', recompute_scale_factor=False, align_corners=False)
+        img1_L_resize = F.interpolate(img1_lab[:,:1,:,:] / 50., size=(imgsize, imgsize), mode='bilinear', recompute_scale_factor=False, align_corners=False)
 
-                        color_vector = colorEncoder(img2_resize)
+        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 = 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_img = torch.cat((img1_lab[:,:1,:,:], fake_ab), 1)
-                        fake_img = Lab2RGB_out(fake_img)
+        fake_img = torch.cat((img1_lab[:,:1,:,:], fake_ab), 1)
+        fake_img = Lab2RGB_out(fake_img)
 
-                        out_folder = os.path.join(args.output_folder, 'color')
-                        mkdirs(out_folder)
-                        out_img_path = os.path.join(out_folder, f'{os.path.splitext(input_file)[0]}_panel_{i}_color.png')
-                        io.imsave(out_img_path, fake_img)
+        out_folder = os.path.dirname(img_path)
+        mkdirs(out_folder)
+        out_img_path = os.path.join(out_folder, f'{img_name}_color.png')
+        io.imsave(out_img_path, fake_img)
 
-    print(f'Colored images have been saved to: {os.path.join(args.output_folder, "color")}')
+    print(f'Colored image has been saved to {out_img_path}.')