tidalove/adain

1import os2import argparse3import torch4import time5import numpy as np6from pathlib import Path7from AdaIN import AdaINNet8from PIL import Image9from torchvision.utils import save_image10from utils import adaptive_instance_normalization, transform,linear_histogram_matching, Range, grid_image11from glob import glob12 13parser = argparse.ArgumentParser()14parser.add_argument('--content_image', type=str, help='Test image file path')15parser.add_argument('--style_image', type=str, required=True, help='Multiple Style image file path, separated by comma')16parser.add_argument('--decoder_weight', type=str, default='decoder.pth', help='Decoder weight file path')17parser.add_argument('--alpha', type=float, default=1.0, choices=[Range(0.0, 1.0)], help='Alpha [0.0, 1.0] controls style transfer level')18parser.add_argument('--interpolation_weights', type=str, help='Weights of interpolate multiple style images')19parser.add_argument('--cuda', action='store_true', help='Use CUDA')20parser.add_argument('--grid_pth', type=str, default=None, help='Specify a grid image path (default=None) if generate a grid image that contains all style transferred images. \21	if use grid mode, provide 4 style images')22parser.add_argument('--color_control', action='store_true', help='Preserve content color')23args = parser.parse_args()24assert args.content_image25assert args.style_image26assert args.decoder_weight27assert args.interpolation_weights or args.grid_pth28 29device = torch.device('cuda' if args.cuda and torch.cuda.is_available() else 'cpu')30 31 32def interpolate_style_transfer(content_tensor, style_tensor, encoder, decoder, alpha=1.0, interpolation_weights=None):33	"""34	Given content image and multiple style images, generate feature maps with encoder, apply 35	neural style transfer with adaptive instance normalization, interpolate style image features 36	with interpolation weights, generate output image with decoder37 38	Args:39		content_tensor (torch.FloatTensor): Content image 40		style_tensor (torch.FloatTensor): Multiple Style Images41		encoder: Encoder (vgg19) network42		decoder: Decoder network43		alpha (float, default=1.0): Weight of style image feature 44		interpolation_weights (list): Weight of each style image 45	46	Return:47		output_tensor (torch.FloatTensor): Interpolate Style Transfer output image48	"""49	50	content_enc = encoder(content_tensor)51	style_enc = encoder(style_tensor)52	53	transfer_enc = torch.zeros_like(content_enc).to(device)54	full_enc = adaptive_instance_normalization(content_enc, style_enc)55	for i, w in enumerate(interpolation_weights):56		transfer_enc += w * full_enc[i]57 58	mix_enc = alpha * transfer_enc + (1-alpha) * content_enc59	return decoder(mix_enc)60	61 62def main():	63	# Read content and style image64	if args.content_image:65		content_pths = [Path(args.content_image)]66	else:67		content_pths = [Path(f) for f in glob(args.content_dir+'/*')]68 69	style_pths_list = args.style_image.split(',')70	style_pths = [Path(pth) for pth in style_pths_list]71 72	assert len(content_pths) > 0, 'Failed to load content image'73	assert len(style_pths) > 0, 'Failed to load style image'74	75	inter_weights = []76	# If grid mode, use 4 style images, 5x5 interpolation weights77	if args.grid_pth:78		assert len(style_pths) == 4, "Under grid mode, specify 4 style images"79		inter_weights = [ [ min(4-a, 4-b) / 4,  min(4-a, b) / 4, min(a, 4-b) / 4, min(a, b) / 4] \80			for a in range(5) for b in range(5) ]81 82	# Use user input interpolation weights83	else:84		inter_weight = [float(i) for i in args.interpolation_weights.split(',')]85		inter_weight = [i / sum(inter_weight) for i in inter_weight]86		inter_weights.append(inter_weight)87	88 89	out_dir = './results_interpolate/'90	os.makedirs(out_dir, exist_ok=True)91	92	# Load AdaIN model93	vgg = torch.load('vgg_normalized.pth')94	model = AdaINNet(vgg).to(device)95	model.decoder.load_state_dict(torch.load(args.decoder_weight))96	model.eval()97	98	# Prepare image transform99	t = transform(512)100 101	imgs = []102 103	for content_pth in content_pths:104		content_tensor = t(Image.open(content_pth)).unsqueeze(0).to(device)105 106		# Prepare multiple style images107		style_tensor = []108		for style_pth in style_pths:109			img = Image.open(style_pth)110			if args.color_control:111				img = transform([512,512])(img).unsqueeze(0)112				img = linear_histogram_matching(content_tensor,img)113				img = img.squeeze(0)114				style_tensor.append(img)115			else:116				style_tensor.append(transform([512, 512])(img))117		style_tensor = torch.stack(style_tensor, dim=0).to(device)118		119		for inter_weight in inter_weights:120			# Execute Interpolate style transfer			121			with torch.no_grad():122				out_tensor = out_tensor = interpolate_style_transfer(content_tensor, style_tensor, model.encoder, model.decoder, args.alpha, inter_weight).cpu()123			124			print("Content: " + content_pth.stem + ". Style: " + str([style_pth.stem for style_pth in style_pths]) + ". Interpolation weight: ", str(inter_weight))125 126			# Save results127			out_pth = out_dir + content_pth.stem + '_interpolate_' + str(inter_weight)128			if args.color_control: out_pth += '_colorcontrol'129			out_pth += content_pth.suffix130			save_image(out_tensor, out_pth)131 132			if args.grid_pth:133				imgs.append(Image.open(out_pth))134 135	# Generate grid image136	if args.grid_pth:137		print("Generating grid image")138		grid_image(5, 5, imgs, save_pth=args.grid_pth)139		print("Finished")140 141if __name__ == '__main__':142	main()