inference.py

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
import argparse

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:,:,:]
    # print(torch.max(img_l), torch.min(img_l))
    # print(torch.max(img_ab), torch.min(img_ab))
    img_l = img_l + 50
    pred_lab = torch.cat((img_l, img_ab), 1)[0,...].numpy()
    # grid_lab = utils.make_grid(pred_lab, nrow=1).numpy().astype("float64")
    # print(grid_lab.shape)
    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):
    # input: rgb, [0, 255], uint8
    img_lab = Normalize(RGB2Lab(inputs))
    img = np.array(inputs, 'float32') # [0, 255]
    img = numpy2tensor(img)
    img_lab = numpy2tensor(img_lab)
    return img.unsqueeze(0), img_lab.unsqueeze(0)

if __name__ == "__main__":
    device = "cuda"

    # model_name = 'Color2Manga_sketch'
    ckpt_path = 'experiments/Color2Manga_gray/074000_gray.pt'
    test_dir_path = 'test_datasets/gray_test'
    no_extractor = False
    # imgs_num = len(os.listdir(test_dir_path)) // 2
    imgsize = 256

    parser = argparse.ArgumentParser()

    parser.add_argument("--path", type=str, default=None, help="path of input image")
    parser.add_argument("--size", type=int, default=None)
    parser.add_argument("--ckpt", type=str, default=None, help="path of model weight")
    parser.add_argument("-ne", "--no_extractor", action='store_true',
                        help="Do not segment the manga panels.")

    args = parser.parse_args()

    if args.path:
        test_dir_path = args.path
    if args.size:
        imgsize = args.size
    if args.ckpt:
        ckpt_path = args.ckpt
    if args.no_extractor:
        no_extractor = args.no_extractor


    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()

    imgs = []
    imgs_lab = []

    # for i in range(imgs_num):
    # idx = i
    # print('Image', idx, 'Input Image', 'in%d.JPEG'%idx, 'Ref Image', 'ref%d.JPEG'%idx)

    while 1:
        print(f'make sure both manga image and reference 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]

        if no_extractor:
            ref_img_path = os.path.join(test_dir_path, input(f"{1}/{1} reference image:"))

            img1 = Image.open(img_path).convert("RGB")
            width, height = img1.size
            img2 = Image.open(ref_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)

            # print('-------',torch.max(img1_lab[:,:1,:,:]), torch.min(img1_lab[:,1:,:,:]))

            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)

                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)
                # io.imsave(out_img_path, 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)

            continue



        # extract panels from manga
        panel_extractor = PanelExtractor(min_pct_panel=5, max_pct_panel=90)
        panels, masks, panel_masks = panel_extractor.extract(img_path)
        panel_num = len(panels)

        ref_img_paths = []
        # ref_img_path = os.path.join(test_dir_path, '%03d_ref.png' % idx)
        print("Please enter the name of the reference image in order according to the number prompts on the picture")
        for i in range(panel_num):
            ref_img_path = os.path.join(test_dir_path, input(f"{i+1}/{panel_num} reference image:"))
            ref_img_paths.append(ref_img_path)




        fake_imgs = []
        for i in range(panel_num):
            img1 = Image.fromarray(panels[i]).convert("RGB")
            width, height = img1.size
            img2 = Image.open(ref_img_paths[i]).convert("RGB")

            # img1 = Image.open(img_path).convert("RGB")
            # width, height = img1.size
            # img2 = Image.open(ref_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)

            # print('-------',torch.max(img1_lab[:,:1,:,:]), torch.min(img1_lab[:,1:,:,:]))

            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)

                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)
                # io.imsave(f'test_datasets/gray_test/panels/{i}.png', fake_img)
                fake_imgs.append(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')

            # show image
            Image.fromarray(fake_imgs[0]).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_imgs[0])
        else:
            panel_extractor.concatPanels(img_path, fake_imgs, masks, panel_masks)

        print(f'colored image has been put to: {test_dir_path}color')