Update pintar.py

pintar.py

@@ -6,7 +6,6 @@ import torch
 import torch.nn.functional as F
 from PIL import Image
 from models import ColorEncoder, ColorUNet
-from extractor.manga_panel_extractor import PanelExtractor
 
 os.environ["CUDA_VISIBLE_DEVICES"] = '0'
 
@@ -54,7 +53,6 @@ if __name__ == "__main__":
     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"
@@ -69,51 +67,43 @@ if __name__ == "__main__":
     colorUNet.load_state_dict(ckpt["colorUNet"])
     colorUNet.eval()
 
-    if args.no_extractor:
-        # Colorize a single image without panel extraction
-        img_path = args.input_folder
-        ref_img_path = args.reference_image
-
-        img1 = Image.open(ref_img_path).convert("RGB")
-        width, height = img1.size
-        img2 = Image.open(img_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, 'colorized_image.png')
-            io.imsave(out_img_path, fake_img)
-
-        if panel_num == 1:
-            out_folder = os.path.dirname(img_path)
-            out_name = os.path.basename(img_path)
-            out_name = os.path.splitext(out_name)[0]
-            out_img_path = os.path.join(out_folder,'color',f'{out_name}_color.png')
-
-            Image.fromarray(fake_imgs[0]).show()
-            folder_path = os.path.join(out_folder, 'color')
-            if not os.path.exists(folder_path):
-                os.makedirs(folder_path)
-            io.imsave(out_img_path, fake_imgs[0])
-        else:
-            panel_extractor.concatPanels(img_path, fake_imgs, masks, panel_masks)
-
-        print(f'Colored images have been saved to: {os.path.join(test_dir_path, "color")}')
+    input_folder = args.input_folder
+    output_folder = args.output_folder
+    reference_image_path = args.reference_image
+
+    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)
+
+                img1 = Image.open(reference_image_path).convert("RGB")
+                width, height = img1.size
+                img2 = Image.open(input_image_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)
+
+                    relative_path = os.path.relpath(input_image_path, input_folder)
+                    output_subfolder = os.path.join(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: {output_folder}')
+