tidalove/adain

1import os2from PIL import Image, ImageFile3import torch4from torch.utils.data import Dataset5import torchvision.transforms as transforms 6import matplotlib.pyplot as plt7from pathlib import Path8from glob import glob9 10def adaptive_instance_normalization(x, y, eps=1e-5):11	"""12	Adaptive Instance Normalization. Perform neural style transfer given content image x13	and style image y.14 15	Args:16		x (torch.FloatTensor): Content image tensor17		y (torch.FloatTensor): Style image tensor18		eps (float, default=1e-5): Small value to avoid zero division19 20	Return:21		output (torch.FloatTensor): AdaIN style transferred output22	"""23 24	mu_x = torch.mean(x, dim=[2, 3])25	mu_y = torch.mean(y, dim=[2, 3])26	mu_x = mu_x.unsqueeze(-1).unsqueeze(-1)27	mu_y = mu_y.unsqueeze(-1).unsqueeze(-1)28 29	sigma_x = torch.std(x, dim=[2, 3])30	sigma_y = torch.std(y, dim=[2, 3])31	sigma_x = sigma_x.unsqueeze(-1).unsqueeze(-1) + eps32	sigma_y = sigma_y.unsqueeze(-1).unsqueeze(-1) + eps33 34	return (x - mu_x) / sigma_x * sigma_y  + mu_y35 36def transform(size):37	"""38	Image preprocess transformation. Resize image and convert to tensor.39 40	Args:41		size (int): Resize image size42 43	Return:44		output (torchvision.transforms): Composition of torchvision.transforms steps45	"""46	47	t = []48	t.append(transforms.Resize(size))49	t.append(transforms.ToTensor())50	t = transforms.Compose(t)51	return t52 53def grid_image(row, col, images, height=6, width=6, save_pth='grid.png'):54	"""55	Generate and save an image that contains row x col grids of images.56 57	Args:58		row (int): number of rows59		col (int): number of columns60		images (list of PIL image): list of images.61		height (int) : height of each image (inch)62		width (int) : width of eac image (inch)63		save_pth (str): save file path64	"""65 66	width = col * width67	height = row * height68	plt.figure(figsize=(width, height))69	for i, image in enumerate(images):70		plt.subplot(row, col, i+1)71		plt.imshow(image)72		plt.axis('off')73		plt.subplots_adjust(wspace=0.01, hspace=0.01)74	plt.savefig(save_pth)75 76 77def linear_histogram_matching(content_tensor, style_tensor):78	"""79	Given content_tensor and style_tensor, transform style_tensor histogram to that of content_tensor.80 81	Args:82		content_tensor (torch.FloatTensor): Content image 83		style_tensor (torch.FloatTensor): Style Image84	85	Return:86		style_tensor (torch.FloatTensor): histogram matched Style Image87	"""88    #for batch89	for b in range(len(content_tensor)):90		std_ct = []91		std_st = []92		mean_ct = []93		mean_st = []94		#for channel95		for c in range(len(content_tensor[b])):96			std_ct.append(torch.var(content_tensor[b][c],unbiased = False))97			mean_ct.append(torch.mean(content_tensor[b][c]))98			std_st.append(torch.var(style_tensor[b][c],unbiased = False))99			mean_st.append(torch.mean(style_tensor[b][c]))100			style_tensor[b][c] = (style_tensor[b][c] - mean_st[c]) * std_ct[c] / std_st[c] + mean_ct[c]101	return style_tensor102 103 104class TrainSet(Dataset):105	"""106	Build Training dataset107	"""108	def __init__(self, content_dir, style_dir, crop_size = 256):109		super().__init__()110 111		self.content_files = [Path(f) for f in glob(content_dir+'/*')]112		self.style_files = [Path(f) for f in glob(style_dir+'/*')]113		114		self.transform = transforms.Compose([115			transforms.Resize(512, interpolation=transforms.InterpolationMode.BICUBIC),116			transforms.RandomCrop(crop_size),117			transforms.ToTensor(),118			transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))119			])120 121		Image.MAX_IMAGE_PIXELS = None122		ImageFile.LOAD_TRUNCATED_IMAGES = True123	124	def __len__(self):125		return min(len(self.style_files), len(self.content_files))126 127	def __getitem__(self, index):128		content_img = Image.open(self.content_files[index]).convert('RGB')129		style_img = Image.open(self.style_files[index]).convert('RGB')130	131		content_sample = self.transform(content_img)132		style_sample = self.transform(style_img)133 134		return content_sample, style_sample135 136class Range(object):137	"""138	Helper class for input argument range restriction139	"""140	def __init__(self, start, end):141		self.start = start142		self.end = end143	def __eq__(self, other):144		return self.start <= other <= self.end