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