readme

README.md

@@ -1,14 +1,15 @@
 2022-AdaIN-pytorch
 ============================
-This is an unofficial Pytorch implementation of the paper, Style Transfer with Adaptive Instance Normalization [arxiv](https://arxiv.org/abs/1703.06868). I referred to the [official implementation](https://github.com/xunhuang1995/AdaIN-style) in Torch. I used pretrained weights for vgg19 and decoder from [naoto0804](https://github.com/naoto0804/pytorch-AdaIN).
+This is an unofficial Pytorch implementation of the paper, 'Style Transfer with Adaptive Instance Normalization' [arxiv](https://arxiv.org/abs/1703.06868). I referred to the [official implementation](https://github.com/xunhuang1995/AdaIN-style) in Torch. I used pretrained weights for vgg19 and decoder from [naoto0804](https://github.com/naoto0804/pytorch-AdaIN).
 
-Requirement
+Requirements
 ----------------------------
-* Python 3.7
+* Python 3.7+
 * PyTorch 1.10
 * Pillow
 * TorchVision
 * Numpy
+* imageio
 * tqdm
 
 
@@ -17,7 +18,7 @@ Usage
 
 ### Training
 
-The encoder uses pretrained vgg19 network. Download the [weight of vgg19](https://drive.google.com/file/d/1UcSl-Zn3byEmn15NIPXMf9zaGCKc2gfx/view?usp=sharing). The decoder is trained on MSCOCO and wikiart dataset. 
+The encoder uses pretrained vgg19 network. Download the [vgg19 weight](https://drive.google.com/file/d/1UcSl-Zn3byEmn15NIPXMf9zaGCKc2gfx/view?usp=sharing). The decoder is trained on MSCOCO and wikiart dataset. 
 Run the script train.py
 ```
 $ python train.py --content_dir $CONTENT_DIR --style_dir STYLE_DIR --cuda
@@ -40,18 +41,13 @@ optional arguments:
 
 ### Test Image Style Transfer
 
-Download the [decoder weight](https://drive.google.com/file/d/18JpLtMOapA-vwBz-LRomyTl24A9GwhTF/view?usp=sharing).
+Download [vgg19 weight](https://drive.google.com/file/d/1UcSl-Zn3byEmn15NIPXMf9zaGCKc2gfx/view?usp=sharing), [decoder weight](https://drive.google.com/file/d/18JpLtMOapA-vwBz-LRomyTl24A9GwhTF/view?usp=sharing) under main folder.
 
-To test basic style transfer, run the script test_image.py.
+To test basic style transfer, run the script test_image.py. Specify '--content_image' or 'style_img' to the image path. Specify '--content_dir' or '--style_dir' to iterate all images under this directory. Specify '--grid_pth' to collect all outputs in a grid image.
 
 ```
 $ python test.py --content_image $IMG --style_image $STYLE --decoder_weight $WEIGHT --cuda
 
-usage: test_style_transfer.py [-h] [--content_image CONTENT_IMAGE] [--content_dir CONTENT_DIR]
-                              [--style_image STYLE_IMAGE] [--style_dir STYLE_DIR]
-                              [--decoder_weight DECODER_WEIGHT] [--alpha {Alpha Range}]
-                              [--cuda] [--grid_pth GRID_PTH]
-
 optional arguments:
   -h, --help            show this help message and exit
   --content_image CONTENT_IMAGE
@@ -67,21 +63,17 @@ optional arguments:
                         Alpha [0.0, 1.0] controls style transfer level
   --cuda                Use CUDA
   --grid_pth GRID_PTH
-                        Specify a grid image path (default=None) if generate a grid image that contains all style transferred images
+                        Specify a grid image path (default=None) if generate a grid image
+                        that contains all style transferred images
 ```
 
 ### Test Image Interpolation Style Transfer
 
-To test style transfer interpolation, run the script test_interpolate.py.
+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. Specify '--grid_pth' to interpolate with different built-in weights.
 
 ```
 $ python test_interpolation.py --content_image $IMG --style_image $STYLE --decoder_weight $WEIGHT --cuda
 
-usage: test_style_transfer.py [-h] [--content_image CONTENT_IMAGE] [--style_image STYLE_IMAGE] 
-                              [--decoder_weight DECODER_WEIGHT] [--alpha {Alpha Range}]
-                              [--interpolation_weights INTERPOLATION_WEIGHTS]
-                              [--cuda] [--grid_pth GRID_PTH]
-
 optional arguments:
   -h, --help            show this help message and exit
   --content_image CONTENT_IMAGE
@@ -92,25 +84,26 @@ optional arguments:
   --alpha {Alpha Range}
                         Alpha [0.0, 1.0] controls style transfer level
   --interpolation_weights INTERPOLATION_WEIGHTS
-                        Interpolation weight of each style image, separated by comma. Perform interpolation style transfer once. Do not specify if input grid_pth.
+                        Interpolation weight of each style image, separated by comma.
+                        Do not specify if input grid_pth.
   --cuda                Use CUDA
   --grid_pth GRID_PTH
-                        Specify a grid image path (default=None) to perform interpolation style transfer multiple times with different built-in weights and generate a grid image that contains all style transferred images. Provide 4 style images. Do not specify if input interpolation_weights.
+                        Specify a grid image path (default=None) to perform interpolation style 
+                        
+                        transfer multiple times with different built-in weights and generate a 
+                        
+                        grid image that contains all style transferred images. Provide 4 style 
+                        
+                        images. Do not specify if input interpolation_weights.
 ```
 
 ### Test Video Style Transfer
-----------------------------
 
 To test video style transfer, run the script test_video.py. 
 
-
 ```
 $ python test_video.py --content_video $VID --style_image $STYLE --decoder_weight $WEIGHT --cuda
 
-usage: test_style_transfer.py [-h] [--content_video CONTENT_VID] [--style_image STYLE_IMAGE]
-                              [--decoder_weight DECODER_WEIGHT] [--alpha {Alpha Range}]
-                              [--cuda]
-
 optional arguments:
   -h, --help            show this help message and exit
   --content_image CONTENT_IMAGE
@@ -124,9 +117,8 @@ optional arguments:
 ```
 
 
-### References
+References
 ----------------------------
-
--[1]: X. Huang and S. Belongie. "Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization.", in ICCV, 2017. [arxiv](https://arxiv.org/abs/1703.06868)
--[2]: [Original implementation in Torch](https://github.com/xunhuang1995/AdaIN-style)
--[3]: [Pretrained weights](https://github.com/naoto0804/pytorch-AdaIN)
+* X. Huang and S. Belongie. "Arbitrary Style Transfer in Real-time with Adaptive Instance Normalization.", in ICCV, 2017. [arxiv](https://arxiv.org/abs/1703.06868)
+* [Original implementation in Torch](https://github.com/xunhuang1995/AdaIN-style)
+* [Pretrained weights](https://github.com/naoto0804/pytorch-AdaIN)

test_style_transfer.py

@@ -1,58 +0,0 @@
-import argparse
-import os
-import numpy as np
-from AdaIN import StyleTransferNet
-from PIL import Image
-import torch
-import torchvision.transforms as transforms 
-from torchvision.utils import save_image
-
-class AlphaRange(object):
-	def __init__(self, start, end):
-		self.start = start
-		self.end = end
-	def __eq__(self, other):
-		return self.start <= other <= self.end
-	def __str__(self):
-		return 'Alpha Range'
-
-parser = argparse.ArgumentParser()
-parser.add_argument('--input_image', type=str, help='test image')
-parser.add_argument('--style_image', type=str, help='style image')
-parser.add_argument('--weight', type=str, help='decoder weight file')
-parser.add_argument('--alpha', type=float, default=1.0, choices=[AlphaRange(0.0, 1.0)], help='Level of style transfer, value between 0 and 1')
-parser.add_argument('--cuda', action='store_true', help='Using GPU to train')
-
-
-if __name__ == '__main__':
-	opt =parser.parse_args()
-	input_image = Image.open(opt.input_image)
-	style_image = Image.open(opt.style_image)
-	output_format = opt.input_image[opt.input_image.find('.'):]
-	out_dir = './results/'
-	os.makedirs(out_dir, exist_ok=True)
-	with torch.no_grad():
-		vgg_model = torch.load('vgg_normalized.pth')
-
-		net = StyleTransferNet(vgg_model)
-		net.decoder.load_state_dict(torch.load(opt.weight))
-
-		net.eval()
-
-		input_image = transforms.Resize(512)(input_image)
-		style_image = transforms.Resize(512)(style_image)
-		
-		input_tensor = transforms.ToTensor()(input_image).unsqueeze(0)
-		style_tensor = transforms.ToTensor()(style_image).unsqueeze(0)
-	
-
-		if torch.cuda.is_available() and opt.cuda:
-			net.cuda()
-			input_tensor = input_tensor.cuda()
-			style_tensor = style_tensor.cuda()
-		out_tensor = net([input_tensor, style_tensor], alpha = opt.alpha)
-
-
-		save_image(out_tensor, out_dir + opt.input_image[opt.input_image.rfind('/')+1: opt.input_image.find('.')]
-								+"_style_"+ opt.style_image[opt.style_image.rfind('/')+1: opt.style_image.find('.')]
-								+ output_format)