Merge pull request #3 from MasaTate/main

.gitignore

@@ -0,0 +1,5 @@
+#Ignore __pycache__
+/__pycache__/
+
+#Ignore results
+/results*/

README.md

@@ -73,7 +73,9 @@ optional arguments:
 To test style transfer interpolation, run the script test_interpolate.py. Specify `--style_image` with multiple paths separated by comma. Specify `--interpolation_weights` to interpolate once. All outputs are saved in `./results_interpolate/`. Specify `--grid_pth` to interpolate with different built-in weights and provide 4 style images.
 
 ```
-$ python test_interpolation.py --content_image $IMG --style_image $STYLE --interpolation_weights $WEIGHTS --cuda
+
+$ python test_interpolate.py --content_image $IMG --style_image $STYLE $WEIGHT --cuda
+
 
 optional arguments:
   -h, --help            show this help message and exit

test.py

@@ -8,7 +8,7 @@ from AdaIN import AdaINNet
 from PIL import Image
 from torchvision.utils import save_image
 from torchvision.transforms import ToPILImage
-from utils import adaptive_instance_normalization, grid_image, transform, Range
+from utils import adaptive_instance_normalization, grid_image, transform,linear_histogram_matching, Range
 from glob import glob
 
 parser = argparse.ArgumentParser()
@@ -20,6 +20,7 @@ parser.add_argument('--decoder_weight', type=str, default='decoder.pth', help='D
 parser.add_argument('--alpha', type=float, default=1.0, choices=[Range(0.0, 1.0)], help='Alpha [0.0, 1.0] controls style transfer level')
 parser.add_argument('--cuda', action='store_true', help='Use CUDA')
 parser.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')
+parser.add_argument('--color_control', action='store_true', help='Preserve content color')
 args = parser.parse_args()
 assert args.content_image or args.content_dir
 assert args.style_image or args.style_dir
@@ -103,6 +104,10 @@ def main():
 			
 			style_tensor = t(Image.open(style_pth)).unsqueeze(0).to(device)
 			
+			# Linear Histogram Matching if needed
+			if args.color_control:
+				style_tensor = linear_histogram_matching(content_tensor,style_tensor)
+
 			# Start time
 			tic = time.perf_counter()
 			
@@ -117,7 +122,9 @@ def main():
 			times.append(toc-tic)
 
 			# Save image
-			out_pth = out_dir + content_pth.stem + '_style_' + style_pth.stem + '_alpha' + str(args.alpha) + content_pth.suffix
+			out_pth = out_dir + content_pth.stem + '_style_' + style_pth.stem + '_alpha' + str(args.alpha)
+			if args.color_control: out_pth += '_colorcontrol'
+			out_pth += content_pth.suffix
 			save_image(out_tensor, out_pth)
 
 			if args.grid_pth:

test_interpolate.py

@@ -7,7 +7,7 @@ from pathlib import Path
 from AdaIN import AdaINNet
 from PIL import Image
 from torchvision.utils import save_image
-from utils import adaptive_instance_normalization, transform, Range, grid_image
+from utils import adaptive_instance_normalization, transform,linear_histogram_matching, Range, grid_image
 from glob import glob
 
 parser = argparse.ArgumentParser()
@@ -19,6 +19,7 @@ parser.add_argument('--interpolation_weights', type=str, help='Weights of interp
 parser.add_argument('--cuda', action='store_true', help='Use CUDA')
 parser.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. \
 	if use grid mode, provide 4 style images')
+parser.add_argument('--color_control', action='store_true', help='Preserve content color')
 args = parser.parse_args()
 assert args.content_image
 assert args.style_image
@@ -106,7 +107,13 @@ def main():
 		style_tensor = []
 		for style_pth in style_pths:
 			img = Image.open(style_pth)
-			style_tensor.append(transform([512, 512])(img))
+			if args.color_control:
+				img = transform([512,512])(img).unsqueeze(0)
+				img = linear_histogram_matching(content_tensor,img)
+				img = img.squeeze(0)
+				style_tensor.append(img)
+			else:
+				style_tensor.append(transform([512, 512])(img))
 		style_tensor = torch.stack(style_tensor, dim=0).to(device)
 		
 		for inter_weight in inter_weights:
@@ -117,7 +124,9 @@ def main():
 			print("Content: " + content_pth.stem + ". Style: " + str([style_pth.stem for style_pth in style_pths]) + ". Interpolation weight: ", str(inter_weight))
 
 			# Save results
-			out_pth = out_dir + content_pth.stem + '_interpolate_' + str(inter_weight) + content_pth.suffix
+			out_pth = out_dir + content_pth.stem + '_interpolate_' + str(inter_weight)
+			if args.color_control: out_pth += '_colorcontrol'
+			out_pth += content_pth.suffix
 			save_image(out_tensor, out_pth)
 
 			if args.grid_pth:

test_video.py

@@ -4,7 +4,7 @@ import torch
 from pathlib import Path
 from AdaIN import AdaINNet
 from PIL import Image
-from utils import transform, adaptive_instance_normalization, Range
+from utils import transform, adaptive_instance_normalization,linear_histogram_matching, Range
 import cv2
 import imageio
 import numpy as np
@@ -17,6 +17,7 @@ parser.add_argument('--style_image', type=str, required=True, help='Style image
 parser.add_argument('--decoder_weight', type=str, default='decoder.pth', help='Decoder weight file path')
 parser.add_argument('--alpha', type=float, default=1.0, choices=[Range(0.0, 1.0)], help='Alpha [0.0, 1.0] controls style transfer level')
 parser.add_argument('--cuda', action='store_true', help='Use CUDA')
+parser.add_argument('--color_control', action='store_true', help='Preserve content color')
 args = parser.parse_args()
 
 device = torch.device('cuda' if args.cuda and torch.cuda.is_available() else 'cpu')
@@ -67,8 +68,10 @@ def main():
 	# Prepare output video writer
 	out_dir = './results_video/'
 	os.makedirs(out_dir, exist_ok=True)
-	out_pth = Path(out_dir + content_video_pth.stem + '_style_' \
-		+ style_image_pth.stem + content_video_pth.suffix)
+	out_pth = out_dir + content_video_pth.stem + '_style_' + style_image_pth.stem
+	if args.color_control: out_pth += '_colorcontrol'
+	out_pth += content_video_pth.suffix
+	out_pth = Path(out_pth)
 	writer = imageio.get_writer(out_pth, mode='I', fps=fps)
 
 	# Load AdaIN model
@@ -82,6 +85,7 @@ def main():
 	style_tensor = t(style_image).unsqueeze(0).to(device)
 
 
+
 	while content_video.isOpened():
 		ret, content_image = content_video.read()
 		# Failed to read a frame
@@ -90,6 +94,10 @@ def main():
 		
 		content_tensor = t(Image.fromarray(content_image)).unsqueeze(0).to(device)
 		
+		# Linear Histogram Matching if needed
+		if args.color_control:
+			style_tensor = linear_histogram_matching(content_tensor,style_tensor)
+
 		with torch.no_grad():
 			out_tensor = style_transfer(content_tensor, style_tensor, model.encoder
 				, model.decoder, args.alpha).cpu().detach().numpy()

utils.py

@@ -74,6 +74,33 @@ def grid_image(row, col, images, height=6, width=6, save_pth='grid.png'):
 	plt.savefig(save_pth)
 
 
+def linear_histogram_matching(content_tensor, style_tensor):
+	"""
+	Given content_tensor and style_tensor, transform style_tensor histogram to that of content_tensor.
+
+	Args:
+		content_tensor (torch.FloatTensor): Content image 
+		style_tensor (torch.FloatTensor): Style Image
+	
+	Return:
+		style_tensor (torch.FloatTensor): histogram matched Style Image
+	"""
+    #for batch
+	for b in range(len(content_tensor)):
+		std_ct = []
+		std_st = []
+		mean_ct = []
+		mean_st = []
+		#for channel
+		for c in range(len(content_tensor[b])):
+			std_ct.append(torch.var(content_tensor[b][c],unbiased = False))
+			mean_ct.append(torch.mean(content_tensor[b][c]))
+			std_st.append(torch.var(style_tensor[b][c],unbiased = False))
+			mean_st.append(torch.mean(style_tensor[b][c]))
+			style_tensor[b][c] = (style_tensor[b][c] - mean_st[c]) * std_ct[c] / std_st[c] + mean_ct[c]
+	return style_tensor
+
+
 class TrainSet(Dataset):
 	"""
 	Build Training dataset