warnings.filterwarnings('ignore')
if __name__ == '__main__':
utils.set_random_seed(args.seed, args.gpu)
if args.color == 'y':
ch = 1
weight = (1.0,)
elif args.color == 'rgb':
ch = 3
weight = (0.29891 * 3, 0.58661 * 3, 0.11448 * 3)
weight = np.array(weight, dtype=np.float32)
weight = weight[:, np.newaxis, np.newaxis]
print('* loading filelist...', end=' ')
filelist = utils.load_filelist(args.dataset_dir, shuffle=True)
valid_num = int(np.ceil(args.validation_rate * len(filelist)))
valid_list, train_list = filelist[:valid_num], filelist[valid_num:]
print('done')
print('* setup model...', end=' ')
if args.model_name is None:
if args.method == 'noise':
model_name = 'anime_style_noise{}'.format(args.noise_level)
elif args.method == 'scale':
model_name = 'anime_style_scale'
elif args.method == 'noise_scale':
model_name = 'anime_style_noise{}_scale'.format(args.noise_level)
model_path = '{}_{}.npz'.format(model_name, args.color)
else:
model_name = args.model_name.rstrip('.npz')
def main():
p = argparse.ArgumentParser(description='Chainer implementation of waifu2x')
p.add_argument('--gpu', '-g', type=int, default=-1)
p.add_argument('--seed', '-s', type=int, default=11)
p.add_argument('--dataset_dir', '-d', required=True)
p.add_argument('--validation_rate', type=float, default=0.05)
p.add_argument('--nr_rate', type=float, default=0.65)
p.add_argument('--chroma_subsampling_rate', type=float, default=0.5)
p.add_argument('--reduce_memory_usage', action='store_true')
p.add_argument('--out_size', type=int, default=64)
p.add_argument('--max_size', type=int, default=256)
p.add_argument('--active_cropping_rate', type=float, default=0.5)
p.add_argument('--active_cropping_tries', type=int, default=10)
p.add_argument('--random_half_rate', type=float, default=0.0)
p.add_argument('--random_color_noise_rate', type=float, default=0.0)
p.add_argument('--random_unsharp_mask_rate', type=float, default=0.0)
p.add_argument('--learning_rate', type=float, default=0.00025)
p.add_argument('--lr_min', type=float, default=0.00001)
p.add_argument('--lr_decay', type=float, default=0.9)
p.add_argument('--lr_decay_interval', type=int, default=5)
p.add_argument('--batch_size', '-b', type=int, default=16)
p.add_argument('--patches', '-p', type=int, default=64)
p.add_argument('--validation_crop_rate', type=float, default=0.5)
p.add_argument('--downsampling_filters', nargs='+', default=['box'])
p.add_argument('--resize_blur_min', type=float, default=0.95)
p.add_argument('--resize_blur_max', type=float, default=1.05)
p.add_argument('--epoch', '-e', type=int, default=50)
p.add_argument('--inner_epoch', type=int, default=4)
p.add_argument('--finetune', '-f', default=None)
p.add_argument('--model_name', default=None)
p.add_argument('--color', '-c', default='rgb',
choices=['y', 'rgb'])
p.add_argument('--arch', '-a', default='VGG7',
choices=['VGG7', '0', 'UpConv7', '1', 'ResNet10', '2', 'UpResNet10', '3'])
p.add_argument('--method', '-m', default='scale',
choices=['noise', 'scale', 'noise_scale'],)
p.add_argument('--noise_level', '-n', type=int, default=1,
choices=[0, 1, 2, 3])
args = p.parse_args()
if args.arch in srcnn.table:
args.arch = srcnn.table[args.arch]
utils.set_random_seed(args.seed, args.gpu)
if args.color == 'y':
ch = 1
weight = (1.0,)
elif args.color == 'rgb':
ch = 3
weight = (0.29891 * 3, 0.58661 * 3, 0.11448 * 3)
weight = np.array(weight, dtype=np.float32)
weight = weight[:, np.newaxis, np.newaxis]
print('* loading filelist...', end=' ')
filelist = utils.load_filelist(args.dataset_dir, shuffle=True)
valid_num = int(np.ceil(args.validation_rate * len(filelist)))
valid_list, train_list = filelist[:valid_num], filelist[valid_num:]
print('done')
print('* setup model...', end=' ')
if args.model_name is None:
if args.method == 'noise':
model_name = 'anime_style_noise{}'.format(args.noise_level)
elif args.method == 'scale':
model_name = 'anime_style_scale'
elif args.method == 'noise_scale':
model_name = 'anime_style_noise{}_scale'.format(args.noise_level)
model_path = '{}_{}.npz'.format(model_name, args.color)
else:
model_name = args.model_name.rstrip('.npz')
model_path = model_name + '.npz'
if not os.path.exists('epoch'):
os.makedirs('epoch')
model = srcnn.archs[args.arch](ch)
if model.offset % model.inner_scale != 0:
raise ValueError('offset %% inner_scale must be 0.')
elif model.inner_scale != 1 and model.inner_scale % 2 != 0:
raise ValueError('inner_scale must be 1 or an even number.')
if args.finetune is not None:
chainer.serializers.load_npz(args.finetune, model)
if args.gpu >= 0:
chainer.backends.cuda.check_cuda_available()
chainer.backends.cuda.get_device(args.gpu).use()
weight = chainer.backends.cuda.to_gpu(weight)
model.to_gpu()
optimizer = optimizers.Adam(alpha=args.learning_rate)
optimizer.setup(model)
print('done')
valid_config = utils.get_config(args, model, train=False)
train_config = utils.get_config(args, model, train=True)
print('* check forward path...', end=' ')
di = train_config.in_size
do = train_config.out_size
dx = model.xp.zeros((args.batch_size, ch, di, di), dtype=np.float32)
dy = model(dx)
if dy.shape[2:] != (do, do):
raise ValueError('Invlid output size\n'
'Expect: {}\n'
'Actual: ({}, {})'.format(dy.shape[2:], do, do))
print('done')
print('* starting processes of dataset sampler...', end=' ')
valid_queue = DatasetSampler(valid_list, valid_config)
train_queue = DatasetSampler(train_list, train_config)
print('done')
best_count = 0
best_score = 0
best_loss = np.inf
for epoch in range(0, args.epoch):
print('### epoch: {} ###'.format(epoch))
train_queue.reload_switch(init=(epoch < args.epoch - 1))
for inner_epoch in range(0, args.inner_epoch):
best_count += 1
print(' # inner epoch: {}'.format(inner_epoch))
start = time.time()
train_loss = train_inner_epoch(
model, weight, optimizer, train_queue, args.batch_size)
if args.reduce_memory_usage:
train_queue.wait()
if train_loss < best_loss:
best_loss = train_loss
print(' * best loss on training dataset: {:.6f}'.format(
train_loss))
valid_score = valid_inner_epoch(
model, valid_queue, args.batch_size)
if valid_score > best_score:
best_count = 0
best_score = valid_score
print(' * best score on validation dataset: PSNR {:.6f} dB'
.format(valid_score))
best_model = model.copy().to_cpu()
epoch_path = 'epoch/{}_epoch{}.npz'.format(model_name, epoch)
chainer.serializers.save_npz(model_path, best_model)
shutil.copy(model_path, epoch_path)
if best_count >= args.lr_decay_interval:
best_count = 0
optimizer.alpha *= args.lr_decay
if optimizer.alpha < args.lr_min:
optimizer.alpha = args.lr_min
else:
print(' * learning rate decay: {:.6f}'.format(
optimizer.alpha))
print(' * elapsed time: {:.6f} sec'.format(time.time() - start))
args = p.parse_args()
if args.arch in srcnn.table:
args.arch = srcnn.table[args.arch]
if __name__ == '__main__':
models = load_models(args)
input_exts = ['.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp']
output_exts = ['.png', '.webp']
outext = '.' + args.extension
outname = None
outdir = args.output
if os.path.isdir(args.input):
filelist = utils.load_filelist(args.input)
else:
tmpname, tmpext = os.path.splitext(os.path.basename(args.output))
if tmpext in output_exts:
outext = tmpext
outname = tmpname
outdir = os.path.dirname(args.output)
outdir = './' if outdir == '' else outdir
elif not tmpext == '':
raise ValueError('Format {} is not supported'.format(tmpext))
filelist = [args.input]
if not os.path.exists(outdir):
os.makedirs(outdir)
for path in filelist:
def main():
p = argparse.ArgumentParser(description='Chainer implementation of waifu2x')
p.add_argument('--gpu', '-g', type=int, default=-1)
p.add_argument('--input', '-i', default='images/small.png')
p.add_argument('--output', '-o', default='./')
p.add_argument('--quality', '-q', type=int, default=None)
p.add_argument('--model_dir', '-d', default=None)
p.add_argument('--scale_ratio', '-s', type=float, default=2.0)
p.add_argument('--tta', '-t', action='store_true')
p.add_argument('--batch_size', '-b', type=int, default=16)
p.add_argument('--block_size', '-l', type=int, default=128)
p.add_argument('--extension', '-e', default='png',
choices=['png', 'webp'])
p.add_argument('--arch', '-a', default='VGG7',
choices=['VGG7', '0', 'UpConv7', '1', 'ResNet10', '2', 'UpResNet10', '3'])
p.add_argument('--method', '-m', default='scale',
choices=['noise', 'scale', 'noise_scale'])
p.add_argument('--noise_level', '-n', type=int, default=1,
choices=[0, 1, 2, 3])
p.add_argument('--color', '-c', default='rgb',
choices=['y', 'rgb'])
p.add_argument('--tta_level', '-T', type=int, default=8,
choices=[2, 4, 8])
g = p.add_mutually_exclusive_group()
g.add_argument('--width', '-W', type=int, default=0)
g.add_argument('--height', '-H', type=int, default=0)
g.add_argument('--shorter_side', '-S', type=int, default=0)
g.add_argument('--longer_side', '-L', type=int, default=0)
args = p.parse_args()
if args.arch in srcnn.table:
args.arch = srcnn.table[args.arch]
models = load_models(args)
input_exts = ['.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp']
output_exts = ['.png', '.webp']
outext = '.' + args.extension
outname = None
outdir = args.output
if os.path.isdir(args.input):
filelist = utils.load_filelist(args.input)
else:
tmpname, tmpext = os.path.splitext(os.path.basename(args.output))
if tmpext in output_exts:
outext = tmpext
outname = tmpname
outdir = os.path.dirname(args.output)
outdir = './' if outdir == '' else outdir
elif tmpext != '':
raise ValueError('Format {} is not supported'.format(tmpext))
filelist = [args.input]
if not os.path.exists(outdir):
os.makedirs(outdir)
for path in filelist:
tmpname, tmpext = os.path.splitext(os.path.basename(path))
if outname is None or len(filelist) > 1:
outname = tmpname
outpath = os.path.join(outdir, '{}{}'.format(outname, outext))
if tmpext.lower() in input_exts:
src = Image.open(path)
w, h = src.size[:2]
if args.width != 0:
args.scale_ratio = args.width / w
elif args.height != 0:
args.scale_ratio = args.height / h
elif args.shorter_side != 0:
if w < h:
args.scale_ratio = args.shorter_side / w
else:
args.scale_ratio = args.shorter_side / h
elif args.longer_side != 0:
if w > h:
args.scale_ratio = args.longer_side / w
else:
args.scale_ratio = args.longer_side / h
dst = src.copy()
start = time.time()
outname += '_(tta{})'.format(args.tta_level) if args.tta else '_'
if 'noise_scale' in models:
outname += '(noise{}_scale{:.1f}x)'.format(
args.noise_level, args.scale_ratio)
dst = upscale_image(
args, dst, models['noise_scale'], models['alpha'])
else:
if 'noise' in models:
outname += '(noise{})'.format(args.noise_level)
dst = denoise_image(args, dst, models['noise'])
if 'scale' in models:
outname += '(scale{:.1f}x)'.format(args.scale_ratio)
dst = upscale_image(args, dst, models['scale'])
print('Elapsed time: {:.6f} sec'.format(time.time() - start))
outname += '({}_{}){}'.format(args.arch, args.color, outext)
if os.path.exists(outpath):
outpath = os.path.join(outdir, outname)
lossless = args.quality is None
quality = 100 if lossless else args.quality
icc_profile = src.info.get('icc_profile')
icc_profile = "" if icc_profile is None else icc_profile
dst.convert(src.mode).save(
outpath, quality=quality, lossless=lossless,
icc_profile=icc_profile)
six.print_('Saved as \'{}\''.format(outpath))
def train():
if args.color == 'y':
ch = 1
weight = (1.0, )
elif args.color == 'rgb':
ch = 3
weight = (0.29891 * 3, 0.58661 * 3, 0.11448 * 3)
weight = np.array(weight, dtype=np.float32)
weight = weight[:, np.newaxis, np.newaxis]
print('* loading filelist...', end=' ')
filelist = utils.load_filelist(args.dataset_dir, shuffle=True)
valid_num = int(np.ceil(args.validation_rate * len(filelist)))
valid_list, train_list = filelist[:valid_num], filelist[valid_num:]
print('done')
print('* loading model...', end=' ')
if args.model_name is None:
if args.method == 'noise':
model_name = ('anime_style_noise%d_%s' %
(args.noise_level, args.color))
elif args.method == 'scale':
model_name = 'anime_style_scale_%s' % args.color
elif args.method == 'noise_scale':
model_name = ('anime_style_noise%d_scale_%s' %
(args.noise_level, args.color))
else:
model_name = args.model_name.rstrip('.npz')
model_path = model_name + '.npz'
if not os.path.exists('epoch'):
os.makedirs('epoch')
model = srcnn.archs[args.arch](ch)
if args.finetune is not None:
chainer.serializers.load_npz(args.finetune, model)
if args.gpu >= 0:
cuda.check_cuda_available()
cuda.get_device(args.gpu).use()
weight = cuda.cupy.array(weight)
model.to_gpu()
optimizer = optimizers.Adam(alpha=args.learning_rate)
optimizer.setup(model)
print('done')
valid_config = get_config(args, model, train=False)
train_config = get_config(args, model, train=True)
print('* starting processes of dataset sampler...', end=' ')
valid_queue = DatasetSampler(valid_list, valid_config)
train_queue = DatasetSampler(train_list, train_config)
print('done')
best_count = 0
best_score = 0
best_loss = np.inf
for epoch in range(0, args.epoch):
print('### epoch: %d ###' % epoch)
train_queue.reload_switch(init=(epoch < args.epoch - 1))
for inner_epoch in range(0, args.inner_epoch):
best_count += 1
print(' # inner epoch: %d' % inner_epoch)
start = time.time()
train_loss = train_inner_epoch(model, weight, optimizer,
train_queue, args.batch_size)
train_queue.wait()
if train_loss < best_loss:
best_loss = train_loss
print(' * best loss on train dataset: %f' % train_loss)
valid_score = valid_inner_epoch(model, valid_queue,
args.batch_size)
if valid_score > best_score:
best_count = 0
best_score = valid_score
print(' * best score on validation dataset: PSNR %f dB' %
valid_score)
best_model = model.copy().to_cpu()
epoch_path = 'epoch/%s_epoch%d.npz' % (model_name, epoch)
chainer.serializers.save_npz(model_path, best_model)
shutil.copy(model_path, epoch_path)
if best_count >= args.lr_decay_interval:
best_count = 0
optimizer.alpha *= args.lr_decay
if optimizer.alpha < args.lr_min:
optimizer.alpha = args.lr_min
else:
print(' * learning rate decay: %f' % optimizer.alpha)
print(' * elapsed time: %f sec' % (time.time() - start))
def main():
p = argparse.ArgumentParser()
p.add_argument('--gpu', '-g', type=int, default=-1)
p.add_argument('--input', '-i', default='../images/original.png')
p.add_argument('--arch', '-a', default='')
p.add_argument('--tta', '-t', action='store_true')
p.add_argument('--batch_size', '-b', type=int, default=16)
p.add_argument('--block_size', '-l', type=int, default=128)
p.add_argument('--chroma_subsampling', '-j', action='store_true')
p.add_argument('--downsampling_filter', '-d', default='box')
p.add_argument('--method',
'-m',
default='scale',
choices=['scale', 'noise_scale'])
p.add_argument('--noise_level', '-n', type=int, default=1, choices=[0, 1])
p.add_argument('--color', '-c', default='rgb', choices=['y', 'rgb'])
p.add_argument('--tta_level', '-T', type=int, default=8, choices=[2, 4, 8])
args = p.parse_args()
if args.arch in srcnn.table:
args.arch = srcnn.table[args.arch]
utils.set_random_seed(0, args.gpu)
if os.path.isdir(args.input):
filelist = utils.load_filelist(args.input)
else:
filelist = [args.input]
images = []
for path in filelist:
basename = os.path.basename(path)
_, ext = os.path.splitext(basename)
if ext.lower() in ['.png', '.bmp', '.tif', '.tiff']:
img = Image.open(path).convert('RGB')
w, h = img.size[:2]
img = img.crop((0, 0, w - (w % 2), h - (h % 2)))
images.append(img)
qualities = [50, 60, 70, 80, 90, 100]
sampling_factor = '1x1,1x1,1x1'
if args.chroma_subsampling:
sampling_factor = '2x2,1x1,1x1'
arch_scores = {}
for arch in srcnn.table.values():
args.arch = arch
models = load_models(args)
scores = []
sems = []
for quality in qualities:
print(arch, quality)
start = time.time()
score, sem = benchmark(args, models, images, sampling_factor,
quality)
scores.append(score)
sems.append(sem)
print('Elapsed time: {:.6f} sec'.format(time.time() - start))
arch_scores[arch] = [scores, sems]
plt.rcParams['xtick.direction'] = 'out'
plt.rcParams['ytick.direction'] = 'out'
plt.rcParams['font.size'] = 12
plt.rcParams['legend.fontsize'] = 12
title = 'scale ({}) '.format(args.downsampling_filter)
title += 'noise{} ({}) '.format(
args.noise_level, sampling_factor) if 'noise' in args.method else ''
title += 'tta{}'.format(args.tta_level) if args.tta else ''
plt.title(title)
plt.xlabel('JPEG quality')
plt.ylabel('PSNR [dB]')
plt.ylim(27.5, 42)
if args.method == 'scale':
plt.xticks([50, 60, 70, 80, 90, 100], [50, 60, 70, 80, 90, 'lossless'])
else:
plt.xticks([50, 60, 70, 80, 90, 100])
plt.yticks([30, 35, 40])
plt.gca().yaxis.set_minor_locator(tick.MultipleLocator(2.5))
if args.method == 'noise_scale':
if args.noise_level == 0:
plt.axvspan(85, 100, color='b', alpha=0.1, lw=0)
elif args.noise_level == 1:
plt.axvspan(65, 90, color='b', alpha=0.1, lw=0)
plt.grid(which='both', color='gray', linestyle='--')
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['top'].set_visible(False)
plt.gca().yaxis.set_ticks_position('left')
plt.gca().xaxis.set_ticks_position('bottom')
for key, value in arch_scores.items():
plt.errorbar(qualities,
value[0],
yerr=value[1],
fmt='o-',
capsize=3,
label=key)
plt.legend(loc='upper left', edgecolor='white')
plt.show()
