Create pintar.py

pintar.py

@@ -0,0 +1,112 @@
+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
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+
+def mkdirs(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+def Lab2RGB_out(img_lab):
+    img_lab = img_lab.detach().cpu()
+    img_l = img_lab[:,:1,:,:]
+    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")
+    return out
+
+def RGB2Lab(inputs):
+    return color.rgb2lab(inputs)
+
+def Normalize(inputs):
+    l = inputs[:, :, 0:1]
+    ab = inputs[:, :, 1:3]
+    l = l - 50
+    lab = np.concatenate((l, ab), 2)
+    return lab.astype('float32')
+
+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__":
+    device = "cuda"
+
+    ckpt_path = 'experiments/Color2Manga_gray/074000_gray.pt'
+    test_dir_path = 'test_datasets/gray_test'
+    no_extractor = False
+    ref_img_path = 'path_to_your_reference_image.jpg'  # Especifica la ruta de tu imagen de referencia aquí
+
+    ckpt = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
+
+    colorEncoder = ColorEncoder().to(device)
+    colorEncoder.load_state_dict(ckpt["colorEncoder"])
+    colorEncoder.eval()
+
+    colorUNet = ColorUNet().to(device)
+    colorUNet.load_state_dict(ckpt["colorUNet"])
+    colorUNet.eval()
+
+    img1 = Image.open(ref_img_path).convert("RGB")
+    width, height = img1.size
+
+    img1, img1_lab = preprocessing(img1)
+    img1 = img1.to(device)
+    img1_lab = img1_lab.to(device)
+
+    while True:
+        print(f'make sure manga images are under this path: {test_dir_path}')
+        img_path = input("please input the name of image needed to be colorized (with file extension): ")
+        img_path = os.path.join(test_dir_path, img_path)
+        img_name = os.path.basename(img_path)
+        img_name = os.path.splitext(img_name)[0]
+
+        img2 = Image.open(img_path).convert("RGB")
+        img2, img2_lab = preprocessing(img2)
+        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.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')
+
+            # show image
+            Image.fromarray(fake_img).show()
+            # save image
+            folder_path = os.path.join(out_folder, 'color')
+            if not os.path.exists(folder_path):
+                os.mkdir(folder_path)
+            io.imsave(out_img_path, fake_img)