tidalove/adain

1import os2import argparse3import torch4import time5import numpy as np6from pathlib import Path7from AdaIN import AdaINNet8from PIL import Image9from torchvision.utils import save_image10from torchvision.transforms import ToPILImage11from utils import adaptive_instance_normalization, grid_image, transform,linear_histogram_matching, Range12from glob import glob13 14parser = argparse.ArgumentParser()15parser.add_argument('--content_image', type=str, help='Content image file path')16parser.add_argument('--content_dir', type=str, help='Content image folder path')17parser.add_argument('--style_image', type=str, help='Style image file path')18parser.add_argument('--style_dir', type=str, help='Content image folder path')19parser.add_argument('--decoder_weight', type=str, default='decoder.pth', help='Decoder weight file path')20parser.add_argument('--alpha', type=float, default=1.0, choices=[Range(0.0, 1.0)], help='Alpha [0.0, 1.0] controls style transfer level')21parser.add_argument('--cuda', action='store_true', help='Use CUDA')22parser.add_argument('--output_dir', type=str, default="results")23parser.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')24parser.add_argument('--color_control', action='store_true', help='Preserve content color')25args = parser.parse_args()26assert args.content_image or args.content_dir27assert args.style_image or args.style_dir28assert args.decoder_weight29 30device = torch.device('cuda' if args.cuda and torch.cuda.is_available() else 'cpu')31 32 33def style_transfer(content_tensor, style_tensor, encoder, decoder, alpha=1.0):34	"""35	Given content image and style image, generate feature maps with encoder, apply 36	neural style transfer with adaptive instance normalization, generate output image37	with decoder38 39	Args:40		content_tensor (torch.FloatTensor): Content image 41		style_tensor (torch.FloatTensor): Style Image42		encoder: Encoder (vgg19) network43		decoder: Decoder network44		alpha (float, default=1.0): Weight of style image feature 45	46	Return:47		output_tensor (torch.FloatTensor): Style Transfer output image48	"""49 50	content_enc = encoder(content_tensor)51	style_enc = encoder(style_tensor)52 53	transfer_enc = adaptive_instance_normalization(content_enc, style_enc)54	55	mix_enc = alpha * transfer_enc + (1-alpha) * content_enc56	return decoder(mix_enc)57 58 59def main():	60	# Read content images and style images61	if args.content_image:62		content_pths = [Path(args.content_image)]63	else:64		content_pths = [Path(f) for f in glob(args.content_dir+'/*')]65	66	if args.style_image:67		style_pths = [Path(args.style_image)]68	else:69		style_pths = [Path(f) for f in glob(args.style_dir+'/*')]70 71	assert len(content_pths) > 0, 'Failed to load content image'72	assert len(style_pths) > 0, 'Failed to load style image'73 74	# Prepare directory for saving results75	os.makedirs(args.output_dir, exist_ok=True)76 77	# Load AdaIN model78	vgg = torch.load('vgg_normalized.pth', weights_only=False)79	model = AdaINNet(vgg).to(device)80	model.decoder.load_state_dict(torch.load(args.decoder_weight, weights_only=False))81	model.eval()82	83	# Prepare image transform84	t = transform(512)85	86	# Prepare grid image, add style images to the first row87	if args.grid_pth:88		# Add empty image89		imgs = [np.ones((1, 1, 3), np.uint8) * 255]90		for style_pth in style_pths:91			imgs.append(Image.open(style_pth))92	93	# Timer94	times = []95 96	for content_pth in content_pths:97		content_img = Image.open(content_pth)98		if not content_img.mode == "RGB":99			content_img = content_img.convert("RGB")100		content_tensor = t(content_img).unsqueeze(0).to(device)101		102		if args.grid_pth:103			imgs.append(content_img)104 105		for style_pth in style_pths:106 107			# check if style transferred image exists already108			out_pth = os.path.join(args.output_dir, content_pth.stem + '_style_' + style_pth.stem + '_alpha' + str(args.alpha) + content_pth.suffix)109			if os.path.isfile(out_pth):110				print("Skipping existing file")111				continue112 113			style_img = Image.open(style_pth)114 115			if not style_img.mode == "RGB":116				style_img = style_img.convert("RGB")117			118			style_tensor = t(style_img).unsqueeze(0).to(device)119			120			# Linear Histogram Matching if needed121			if args.color_control:122				style_tensor = linear_histogram_matching(content_tensor,style_tensor)123 124			# Start time125			tic = time.perf_counter()126			127			# Execute style transfer128			with torch.no_grad():129				out_tensor = style_transfer(content_tensor, style_tensor, model.encoder, model.decoder, args.alpha).cpu()130		131			 # End time132			toc = time.perf_counter()133			print("Content: " + content_pth.stem + ". Style: " \134				+ style_pth.stem + '. Alpha: ' + str(args.alpha) + '. Style Transfer time: %.4f seconds' % (toc-tic))135			times.append(toc-tic)136 137			# Save image138			save_image(out_tensor, out_pth)139 140			if args.grid_pth:141				imgs.append(Image.open(out_pth))	142 143	# Remove runtime of first iteration because it is flawed for some unknown reason144	if len(times) > 1:145		times.pop(0)146		avg = sum(times)/len(times)147		print("Average style transfer time: %.4f seconds" % (avg))148 149	# Generate grid image150	if args.grid_pth:151		print("Generating grid image")152		grid_image(len(content_pths) + 1, len(style_pths) + 1, imgs, save_pth=args.grid_pth)153		print("Finished")154 155if __name__ == '__main__':156	main()