opt.load_model = True
opt.num_threads = 1 # test code only supports num_threads = 1
opt.batch_size = 1 # test code only supports batch_size = 1
opt.display_id = -1 # no visdom display
opt.phase = 'val'
opt.dataroot = './dataset/ilsvrc2012/%s/' % opt.phase
opt.serial_batches = True
opt.aspect_ratio = 1.
# process opt.suffix
if opt.suffix:
suffix = ('_' + opt.suffix.format(**vars(opt))) if opt.suffix != '' else ''
opt.name = opt.name + suffix
opt.A = 2 * opt.ab_max / opt.ab_quant + 1
opt.B = opt.A
class Colorization(torch.nn.Module):
def __init__(self):
super(Colorization, self).__init__()
model = create_model(opt)
model.setup(opt)
model.eval()
self.model = model
def forward(self, image, hint):
data = {
"A": image[:, 0:1, :, :],
"B": image[:, 1:3, :, :],
"hint_B": hint[:, 0:2, :, :],
