Update pintar.py

pintar.py

@@ -36,23 +36,13 @@ 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.")
     args = parser.parse_args()
 
     device = "cuda"
@@ -68,29 +58,35 @@ if __name__ == "__main__":
     colorUNet.eval()
 
     reference_img = Image.open(args.reference_image).convert("RGB")
-    reference_img, reference_img_lab = preprocessing(reference_img)
-    reference_img = reference_img.to(device)
-    reference_img_lab = reference_img_lab.to(device)
+    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, img_lab = preprocessing(Image.open(input_image_path).convert("RGB"))
-                img = img.to(device)
-                img_lab = img_lab.to(device)
+                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)
 
                 with torch.no_grad():
-                    img_resize = F.interpolate(img / 255., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
-                    img_L_resize = F.interpolate(img_lab[:, :1, :, :] / 50., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
+                    img_resize = F.interpolate(img_lab / 110., size=(256, 256), mode='bilinear', recompute_scale_factor=False, align_corners=False)
+                    img_L_resize = F.interpolate(img_resize[:, :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 * 110, size=(img.size(2), img.size(3)), mode='bilinear', recompute_scale_factor=False, align_corners=False)
+                    fake_ab = F.interpolate(fake_ab, size=(img.shape[0], img.shape[1]), mode='bilinear', recompute_scale_factor=False, align_corners=False)
 
                     fake_img = torch.cat((img_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')