import torch import torch.nn as nn import numpy as np from utils.dataset_utils import get_sketch from utils.utils import resize_pad, generate_mask, extract_cbr, create_cbz, sorted_alphanumeric, subfolder_image_search, remove_folder from torchvision.transforms import ToTensor import os import matplotlib.pyplot as plt import argparse from model.models import Colorizer, Generator from model.extractor import get_seresnext_extractor from utils.xdog import XDoGSketcher from utils.utils import open_json import sys from denoising.denoiser import FFDNetDenoiser def colorize_without_hint(inp, color_args): i_hint = torch.zeros(1, 4, inp.shape[2], inp.shape[3]).float().to(color_args['device']) with torch.no_grad(): fake_color, _ = color_args['colorizer'](torch.cat([inp, i_hint], 1)) if color_args['auto_hint']: mask = generate_mask(fake_color.shape[2], fake_color.shape[3], full = False, prob = 1, sigma = color_args['auto_hint_sigma']).unsqueeze(0) mask = mask.to(color_args['device']) if color_args['ignore_gray']: diff1 = torch.abs(fake_color[:, 0] - fake_color[:, 1]) diff2 = torch.abs(fake_color[:, 0] - fake_color[:, 2]) diff3 = torch.abs(fake_color[:, 1] - fake_color[:, 2]) mask = ((mask + ((diff1 + diff2 + diff3) > 60 / 255).float().unsqueeze(1)) == 2).float() i_hint = torch.cat([fake_color * mask, mask], 1) with torch.no_grad(): fake_color, _ = color_args['colorizer'](torch.cat([inp, i_hint], 1)) return fake_color def process_image(image, color_args, to_tensor = ToTensor()): image, pad = resize_pad(image) if color_args['denoiser'] is not None: image = color_args['denoiser'].get_denoised_image(image, color_args['denoiser_sigma']) bw, dfm = get_sketch(image, color_args['sketcher'], color_args['dfm']) bw = to_tensor(bw).unsqueeze(0).to(color_args['device']) dfm = to_tensor(dfm).unsqueeze(0).to(color_args['device']) output = colorize_without_hint(torch.cat([bw, dfm], 1), color_args) result = output[0].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5 if pad[0] != 0: result = result[:-pad[0]] if pad[1] != 0: result = result[:, :-pad[1]] return result def colorize_with_hint(inp, color_args): with torch.no_grad(): fake_color, _ = color_args['colorizer'](inp) return fake_color def process_image_with_hint(bw, dfm, hint, color_args, to_tensor = ToTensor()): bw = to_tensor(bw).unsqueeze(0).to(color_args['device']) dfm = to_tensor(dfm).unsqueeze(0).to(color_args['device']) i_hint = (torch.FloatTensor(hint[..., :3]).permute(2, 0, 1) - 0.5) / 0.5 mask = torch.FloatTensor(hint[..., 3:]).permute(2, 0, 1) i_hint = torch.cat([i_hint * mask, mask], 0).unsqueeze(0).to(color_args['device']) output = colorize_with_hint(torch.cat([bw, dfm, i_hint], 1), color_args) result = output[0].cpu().permute(1, 2, 0).numpy() * 0.5 + 0.5 return result def colorize_single_image(file_path, save_path, color_args): try: image = plt.imread(file_path) colorization = process_image(image, color_args) plt.imsave(save_path, colorization) return True except KeyboardInterrupt: sys.exit(0) except: print('Failed to colorize {}'.format(file_path)) return False def colorize_images(source_path, target_path, color_args): images = os.listdir(source_path) for image_name in images: file_path = os.path.join(source_path, image_name) name, ext = os.path.splitext(image_name) if (ext != '.png'): image_name = name + '.png' save_path = os.path.join(target_path, image_name) colorize_single_image(file_path, save_path, color_args) def colorize_cbr(file_path, color_args): file_name = os.path.splitext(os.path.basename(file_path))[0] temp_path = 'temp_colorization' if not os.path.exists(temp_path): os.makedirs(temp_path) extract_cbr(file_path, temp_path) images = subfolder_image_search(temp_path) result_images = [] for image_path in images: save_path = image_path path, ext = os.path.splitext(save_path) if (ext != '.png'): save_path = path + '.png' res_flag = colorize_single_image(image_path, save_path, color_args) result_images.append(save_path if res_flag else image_path) result_name = os.path.join(os.path.dirname(file_path), file_name + '_colorized.cbz') create_cbz(result_name, result_images) remove_folder(temp_path) return result_name def parse_args(): parser = argparse.ArgumentParser() parser.add_argument("-p", "--path", required=True) parser.add_argument("-gen", "--generator", default = 'model/generator.pth') parser.add_argument("-ext", "--extractor", default = 'model/extractor.pth') parser.add_argument("-s", "--sigma", type = float, default = 0.003) parser.add_argument('-g', '--gpu', dest = 'gpu', action = 'store_true') parser.add_argument('-ah', '--auto', dest = 'autohint', action = 'store_true') parser.add_argument('-ig', '--ignore_grey', dest = 'ignore', action = 'store_true') parser.add_argument('-nd', '--no_denoise', dest = 'denoiser', action = 'store_false') parser.add_argument("-ds", "--denoiser_sigma", type = int, default = 25) parser.set_defaults(gpu = False) parser.set_defaults(autohint = False) parser.set_defaults(ignore = False) parser.set_defaults(denoiser = True) args = parser.parse_args() return args if __name__ == "__main__": args = parse_args() if args.gpu: device = 'cuda' else: device = 'cpu' generator = Generator() generator.load_state_dict(torch.load(args.generator)) extractor = get_seresnext_extractor() extractor.load_state_dict(torch.load(args.extractor)) colorizer = Colorizer(generator, extractor) colorizer = colorizer.eval().to(device) sketcher = XDoGSketcher() xdog_config = open_json('configs/xdog_config.json') for key in xdog_config.keys(): if key in sketcher.params: sketcher.params[key] = xdog_config[key] denoiser = None if args.denoiser: denoiser = FFDNetDenoiser(device, args.denoiser_sigma) color_args = {'colorizer':colorizer, 'sketcher':sketcher, 'auto_hint':args.autohint, 'auto_hint_sigma':args.sigma,\ 'ignore_gray':args.ignore, 'device':device, 'dfm' : True, 'denoiser':denoiser, 'denoiser_sigma' : args.denoiser_sigma} if os.path.isdir(args.path): colorization_path = os.path.join(args.path, 'colorization') if not os.path.exists(colorization_path): os.makedirs(colorization_path) colorize_images(args.path, colorization_path, color_args) elif os.path.isfile(args.path): split = os.path.splitext(args.path) if split[1].lower() in ('.cbr', '.cbz', '.rar', '.zip'): colorize_cbr(args.path, color_args) elif split[1].lower() in ('.jpg', '.png', ',jpeg'): new_image_path = split[0] + '_colorized' + '.png' colorize_single_image(args.path, new_image_path, color_args) else: print('Wrong format') else: print('Wrong path')