mangaaa / inference.py
Keiser41's picture
Upload 246 files
212d7be
raw
history blame
7.65 kB
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')