Update pintar.py

pintar.py

@@ -1,11 +1,12 @@
 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
+import argparse
 
 os.environ["CUDA_VISIBLE_DEVICES"] = '0'
 
@@ -29,8 +30,6 @@ def Normalize(inputs):
     l = inputs[:, :, 0:1]
     ab = inputs[:, :, 1:3]
     l = l - 50
-    l = l / 50  # Normalizar L al rango [-1, 1]
-    ab = ab / 110  # Normalizar ab al rango [-1, 1]
     lab = np.concatenate((l, ab), 2)
     return lab.astype('float32')
 
@@ -38,10 +37,20 @@ def numpy2tensor(inputs):
     out = torch.from_numpy(inputs.transpose(2,0,1))
     return out
 
+def tensor2numpy(inputs):
+    out = inputs[0,...].detach().cpu().numpy().transpose(1,2,0)
+    return out
+
+def preprocessing(inputs):
+    img_lab = Normalize(RGB2Lab(inputs))
+    img = np.array(inputs, 'float32')
+    img = numpy2tensor(img)
+    img_lab = numpy2tensor(img_lab)
+    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("-r", "--reference_image", type=str, required=True, help="Path to the reference image for colorization.")
     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.")
@@ -59,41 +68,82 @@ if __name__ == "__main__":
     colorUNet.load_state_dict(ckpt["colorUNet"])
     colorUNet.eval()
 
-    reference_img = Image.open(args.reference_image).convert("RGB")
-    reference_img = np.array(reference_img).astype(np.float32) / 255.0  # Asegúrate de que la referencia esté en el rango [0, 1]
-    reference_img_lab = RGB2Lab(reference_img)
-    reference_img_lab = Normalize(reference_img_lab)
-    reference_img_lab = numpy2tensor(reference_img_lab)
-    reference_img_lab = reference_img_lab.to(device).unsqueeze(0)
-
-    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)
-
-                img = Image.open(input_image_path).convert("RGB")
-                img = np.array(img).astype(np.float32) / 255.0  # Asegúrate de que la imagen de entrada esté en el rango [0, 1]
-                img_lab = RGB2Lab(img)
-                img_lab = Normalize(img_lab)
-                img_lab = numpy2tensor(img_lab)
-                img_lab = img_lab.to(device).unsqueeze(0)
+    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():
-                    img_resize = F.interpolate(img_lab, size=(256, 256), mode='bilinear', align_corners=False)
-                    img_L_resize = F.interpolate(img_resize[:, :1, :, :], size=(256, 256), mode='bilinear', align_corners=False)
+                    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(img_resize)
-                    fake_ab = colorUNet((img_L_resize, color_vector))
-                    fake_ab = F.interpolate(fake_ab, size=(img.shape[0], img.shape[1]), mode='bilinear', 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((img_lab[:, :1, :, :], fake_ab), 1)
+                    fake_img = torch.cat((img1_lab[:, :1, :, :], fake_ab), 1)
                     fake_img = Lab2RGB_out(fake_img)
-                    fake_img = (fake_img * 255).astype(np.uint8)  # Convierte de nuevo a [0, 255]
-
-                    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: {args.output_folder}')
+                    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)
+
+                    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]}_panel_{i}_color.png')
+                        io.imsave(out_img_path, fake_img)
+
+    print(f'Colored images have been saved to: {os.path.join(args.output_folder, "color")}')