def main():
# Prepare WCT model
vgg1 = 'models/vgg_normalised_conv1_1.pth'
vgg2 = 'models/vgg_normalised_conv2_1.pth'
vgg3 = 'models/vgg_normalised_conv3_1.pth'
vgg4 = 'models/vgg_normalised_conv4_1.pth'
vgg5 = 'models/vgg_normalised_conv5_1.pth'
decoder1 = 'models/feature_invertor_conv1_1.pth'
decoder2 = 'models/feature_invertor_conv2_1.pth'
decoder3 = 'models/feature_invertor_conv3_1.pth'
decoder4 = 'models/feature_invertor_conv4_1.pth'
decoder5 = 'models/feature_invertor_conv5_1.pth'
paths = vgg1, vgg2, vgg3, vgg4, vgg5, decoder1, decoder2, decoder3, decoder4, decoder5
wct = WCT(paths)
# Prepare images
content_image = Image.open(args.content).resize((args.content_w, args.content_h))
contentImg = TF.to_tensor(content_image)
contentImg.unsqueeze_(0)
style_image = Image.open(args.style).resize((args.style_w, args.style_h))
styleImg = TF.to_tensor(style_image)
styleImg.unsqueeze_(0)
csF = torch.Tensor()
cImg = Variable(contentImg, volatile=True)
sImg = Variable(styleImg, volatile=True)
csF = Variable(csF)
cImg = cImg.cuda(0)
sImg = sImg.cuda(0)
csF = csF.cuda(0)
wct.cuda(0)
# Run style transfer
start_time = time.time()
styleTransfer(wct, args.alpha, cImg, sImg, csF, args.output)
end_time = time.time()
print('Elapsed time is: %f' % (end_time - start_time))
def style_transfer(args, content_img, style_img, imname, csF):
wct = WCT(args)
if (args.cuda):
wct = wct.cuda(args.gpu)
sF5 = wct.e5(style_img)
cF5 = wct.e5(content_img)
sF5 = sF5.data.cpu().squeeze(0)
cF5 = cF5.data.cpu().squeeze(0)
csF5 = wct.transform(cF5, sF5, csF, args.alpha)
Im5 = wct.d5(csF5)
sF4 = wct.e4(style_img)
cF4 = wct.e4(Im5)
sF4 = sF4.data.cpu().squeeze(0)
cF4 = cF4.data.cpu().squeeze(0)
csF4 = wct.transform(cF4, sF4, csF, args.alpha)
Im4 = wct.d4(csF4)
sF3 = wct.e3(style_img)
cF3 = wct.e3(Im4)
sF3 = sF3.data.cpu().squeeze(0)
cF3 = cF3.data.cpu().squeeze(0)
csF3 = wct.transform(cF3, sF3, csF, args.alpha)
Im3 = wct.d3(csF3)
sF2 = wct.e2(style_img)
cF2 = wct.e2(Im3)
sF2 = sF2.data.cpu().squeeze(0)
cF2 = cF2.data.cpu().squeeze(0)
csF2 = wct.transform(cF2, sF2, csF, args.alpha)
Im2 = wct.d2(csF2)
sF1 = wct.e1(style_img)
cF1 = wct.e1(Im2)
sF1 = sF1.data.cpu().squeeze(0)
cF1 = cF1.data.cpu().squeeze(0)
csF1 = wct.transform(cF1, sF1, csF, args.alpha)
Im1 = wct.d1(csF1)
# save_image has this wired design to pad images with 4 pixels at default.
vutils.save_image(Im1.data.cpu().float(),
os.path.join(args.output_dir, imname))
return
