Update pintar.py

pintar.py

@@ -1,8 +1,10 @@
 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
@@ -20,7 +22,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,19 +51,20 @@ 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", 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")
-
+    parser = argparse.ArgumentParser(description="Colorize manga images based on a single reference image.")
+    parser.add_argument("-i", "--input_folder", type=str, required=True, help="Path to the input folder containing images to be colorized.")
+    parser.add_argument("-r", "--reference_image", type=str, required=True, help="Path to the reference image.")
+    parser.add_argument("-c", "--ckpt", 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 to save colorized images.")
     args = parser.parse_args()
 
     device = "cuda"
-    input_image_dir = args.input
-    output_directory = args.output
-    ckpt_path = args.checkpoint
-    reference_image_path = args.reference
+
+    input_folder = args.input_folder
+    reference_image_path = args.reference_image
+    ckpt_path = args.ckpt
+    output_folder = args.output_folder
+
     imgsize = 256
 
     ckpt = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
@@ -74,32 +77,39 @@ if __name__ == "__main__":
     colorUNet.load_state_dict(ckpt["colorUNet"])
     colorUNet.eval()
 
-    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"))
+    # Recorre recursivamente el directorio de entrada y procesa cada imagen encontrada
+    for root, dirs, files in os.walk(input_folder):
+        for file in files:
+            if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
+                input_image_path = os.path.join(root, file)
+                img_name = os.path.splitext(os.path.basename(input_image_path))[0]
+
+                img1 = Image.open(input_image_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=(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)
+                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.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)
+                    out_subfolder = os.path.join(output_folder, os.path.relpath(root, input_folder))
+                    out_folder = os.path.join(out_subfolder, 'color')
+                    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 image has been saved to {out_img_path}.')
+    print(f'Colored images have been saved to {output_folder}.')