def main():
parser = argparse.ArgumentParser(
description='chainer line drawing colorization')
parser.add_argument('--batchsize',
'-b',
type=int,
default=16,
help='Number of images in each mini-batch')
parser.add_argument('--epoch',
'-e',
type=int,
default=500,
help='Number of sweeps over the dataset to train')
parser.add_argument('--gpu',
'-g',
type=int,
default=-1,
help='GPU ID (negative value indicates CPU)')
parser.add_argument('--dataset',
'-i',
default='./images/',
help='Directory of image files.')
parser.add_argument('--out',
'-o',
default='./result',
help='Directory to output the result')
parser.add_argument('--resume',
'-r',
default='',
help='Resume the training from snapshot')
parser.add_argument('--seed', type=int, default=0, help='Random seed')
parser.add_argument('--snapshot_interval',
type=int,
default=1000,
help='Interval of snapshot')
parser.add_argument('--display_interval',
type=int,
default=10,
help='Interval of displaying log to console')
parser.add_argument('--test_visual_interval',
type=int,
default=500,
help='Interval of drawing test images')
parser.add_argument('--test_out',
default='./test_result/',
help='DIrectory to output test samples')
parser.add_argument('--test_image_path',
default='./test_samples/test_sample3/',
help='Directory of image files for testing')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}'.format(args.batchsize))
print('# epoch: {}'.format(args.epoch))
print('')
root = args.dataset
gen = generator.GEN()
#serializers.load_npz("/mnt/sakura201/gao/lineart/result_cnn/gen_iter_16000", gen)
#print('generator loaded')
dis = discriminator.DIS()
# paired dataset
paired_dataset = Rough2LineDatasetNote("dat/paired_dataset.dat",
root + "rough/",
root + "line/",
root + "note",
train=True,
size=256)
paired_iter = chainer.iterators.MultiprocessIterator(
paired_dataset, args.batchsize)
if args.gpu >= 0:
chainer.cuda.get_device(args.gpu).use() # Make a specified GPU current
gen.to_gpu() # Copy the model to the GPU
dis.to_gpu() # Copy the model to the GPU
# Setup optimizer parameters.
opt = optimizers.Adam(alpha=0.0001)
opt.setup(gen)
opt.add_hook(chainer.optimizer.WeightDecay(1e-5), 'hook_gen')
opt_d = chainer.optimizers.Adam(alpha=0.0001)
opt_d.setup(dis)
opt_d.add_hook(chainer.optimizer.WeightDecay(1e-5), 'hook_dec')
# Set up a trainer
updater = ganUpdater(
models=(gen, dis),
iterator={
'main': paired_iter,
#'test': test_iter
},
optimizer={
'gen': opt,
'dis': opt_d
},
device=args.gpu)
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=args.out)
snapshot_interval = (args.snapshot_interval, 'iteration')
snapshot_interval2 = (args.snapshot_interval * 2, 'iteration')
trainer.extend(extensions.dump_graph('gen/loss'))
trainer.extend(extensions.snapshot(), trigger=snapshot_interval2)
trainer.extend(extensions.snapshot_object(gen,
'gen_iter_{.updater.iteration}'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
dis, 'gen_dis_iter_{.updater.iteration}'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(opt, 'optimizer_'),
trigger=snapshot_interval)
trainer.extend(extensions.LogReport(trigger=(10, 'iteration'), ))
trainer.extend(
extensions.PrintReport([
'epoch', 'gen/loss', 'gen/loss_L', 'gen/loss_adv', 'dis/loss',
'dis/fake_p', 'dis/real_p'
]))
trainer.extend(extensions.ProgressBar(update_interval=10))
trainer.extend(test_samples_simplification(updater, gen, args.test_out,
args.test_image_path),
trigger=(args.test_visual_interval, 'iteration'))
trainer.run()
if args.resume:
# Resume from a snapshot
chainer.serializers.load_npz(args.resume, trainer)
# Save the trained model
chainer.serializers.save_npz(os.path.join(out, 'model_final'), gen)
chainer.serializers.save_npz(os.path.join(out, 'optimizer_final'), opt)
def simplify(file_path, simplifier, output_path, s_size):
origin, transposed = read_img(file_path, s_size)
x_in = np.zeros((1, 1, origin.shape[1], origin.shape[2]), dtype='f')
x_in[0, :] = origin[0]
x_in = cuda.to_gpu(x_in)
cnn_in = Variable(x_in)
cnn_out = simplifier(cnn_in, test=True)
save_as_img(cnn_out.data[0], output_path, origin, transposed)
if __name__ == '__main__':
if not os.path.exists(args.output):
os.makedirs(args.output)
simplifier = generator.GEN()
serializers.load_npz(args.model, simplifier)
print('simplifier loaded')
if args.gpu >= 0:
cuda.get_device(args.gpu).use()
simplifier.to_gpu()
file_test = glob.glob(args.test_path + '/*.jpg')
for f in file_test:
filename = os.path.basename(f)
print(filename)
filename = os.path.splitext(filename)[0]
simplify(f, simplifier, args.output + "/" + filename + ".jpg",
args.size)
x_in = cuda.to_gpu(x_in)
cnn_in = Variable(x_in)
cnn_out = simplifier1(cnn_in, test=True)
if args.mode == 1:
cnn_in = (cnn_out.data * 255).clip(0, 255).astype(np.float32)
cnn_out = simplifier2(cnn_in, test=True)
save_as_img(cnn_out.data[0], output_path, origin, transposed)
if __name__ == '__main__':
if not os.path.exists(args.output):
os.makedirs(args.output)
simplifier1 = generator.GEN()
serializers.load_npz(args.model1, simplifier1)
simplifier2 = generator.GEN()
serializers.load_npz(args.model2, simplifier2)
print('simplifier loaded')
if args.gpu >= 0:
cuda.get_device(args.gpu).use()
simplifier1.to_gpu()
simplifier2.to_gpu()
file_test = glob.glob(args.test_path + '/*.jpg')
for f in file_test:
filename = os.path.basename(f)
print(filename)
filename = os.path.splitext(filename)[0]