add EFDM

.gitattributes

@@ -25,3 +25,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pth.tar filter=lfs diff=lfs merge=lfs -text
+examples/** filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,130 @@
+import gradio as gr
+import toml
+import torch
+from PIL import Image
+from torch import nn
+from torchvision import transforms
+
+import net
+from function import *
+
+cfg = toml.load("config.toml")  # static variables
+
+# Setup device
+if torch.cuda.is_available() and cfg["use_cuda"]:
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
+# Load pretrained models
+decoder = net.decoder
+vgg = net.vgg
+
+decoder.eval()
+vgg.eval()
+
+decoder.load_state_dict(torch.load(cfg["decoder_weight"]))
+vgg.load_state_dict(torch.load(cfg["vgg_weight"]))
+vgg = nn.Sequential(*list(vgg.children())[:31])
+
+vgg = vgg.to(device)
+decoder = decoder.to(device)
+
+
+def transform(img, size, crop):
+    transform_list = []
+    if size > 0:
+        transform_list.append(transforms.Resize(size))
+    if crop:
+        transform_list.append(transforms.CenterCrop(size))
+    transform_list.append(transforms.ToTensor())
+    transform = transforms.Compose(transform_list)
+    return transform(img)
+
+
+@torch.inference_mode()
+def style_transfer(content, style, style_type, alpha, keep_resolution):
+    """Stylize function"""
+    style_type = style_type.lower()
+
+    # Step 1: convert image to PyTorch Tensor
+    if keep_resolution:
+        style = style.resize(content.size, Image.ANTIALIAS)
+
+    if style_type == "efdm" and not keep_resolution:
+        content = transform(content, cfg["content_size"], cfg["crop"])
+        style = transform(style, cfg["style_size"], cfg["crop"])
+    else:
+        content = transform(content, -1, False)
+        style = transform(style, -1, False)
+
+    content = content.to(device).unsqueeze(0)
+    style = style.to(device).unsqueeze(0)
+
+    # Step 2: extract content feature and style feature
+    content_feat = vgg(content)
+    style_feat = vgg(style)
+
+    # Step 3: perform style transfer
+    transfer = {
+        "adain": adaptive_instance_normalization,
+        "adamean": adaptive_mean_normalization,
+        "adastd": adaptive_std_normalization,
+        "efdm": exact_feature_distribution_matching,
+        "hm": histogram_matching,
+    }[style_type]
+    feat = transfer(content_feat, style_feat)
+
+    # Step 4: content-style trade-off
+    feat = feat * alpha + content_feat * (1 - alpha)
+
+    # Step 5: decode to image
+    output = decoder(feat).cpu().squeeze(0).clamp_(0, 1)
+    output = transforms.ToPILImage()(output)
+
+    torch.cuda.ipc_collect()
+    torch.cuda.empty_cache()
+
+    return output
+
+
+# Add image examples
+example_img_pairs = {
+    "examples/content/sailboat.jpg": "examples/style/sketch.png",
+    "examples/content/granatum.jpg": "examples/style/flowers_in_a_turquoise_vase.jpg",
+    "examples/content/einstein.jpeg": "examples/style/polasticot2.jpeg",
+    "examples/content/paris.jpeg": "examples/style/vangogh.jpeg",
+}
+
+# Customize interface
+title = "Style Transfer with EFDM"
+description = """
+Gradio demo for neural style transfer using exact feature distribution matching
+"""
+article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2203.07740'>Exact Feature Distribution Matching for Arbitrary Style Transfer and Domain Generalization</a></p>"
+content_input = gr.inputs.Image(label="Content Image", source="upload", type="pil")
+style_input = gr.inputs.Image(label="Style Image", source="upload", type="pil")
+style_type = gr.inputs.Radio(
+    ["EFDM", "AdaIN", "AdaMean", "AdaStd", "HM"], label="Method"
+)
+alpha_selector = gr.inputs.Slider(
+    minimum=0.0, maximum=1.0, step=0.01, default=1.0, label="Content-Style trade-off"
+)
+keep_resolution = gr.inputs.Checkbox(
+    default=False, label="Keep content image resolution"
+)
+
+iface = gr.Interface(
+    fn=style_transfer,
+    inputs=[content_input, style_input, style_type, alpha_selector, keep_resolution],
+    outputs=["image"],
+    title=title,
+    description=description,
+    article=article,
+    theme="huggingface",
+    examples=[
+        [content, style, "EFDM", 1.0, False]
+        for content, style in example_img_pairs.items()
+    ],
+)
+iface.launch(debug=False, enable_queue=True)

config.toml

@@ -0,0 +1,8 @@
+use_cuda = true
+
+content_size = 512
+style_size = 512
+crop = true
+
+vgg_weight = "pretrained/vgg_normalised.pth"
+decoder_weight = "pretrained/efdm_decoder_iter_160000.pth.tar"

function.py

@@ -0,0 +1,112 @@
+import torch
+from skimage.exposure import match_histograms
+import numpy as np
+
+def calc_mean_std(feat, eps=1e-5):
+    # eps is a small value added to the variance to avoid divide-by-zero.
+    size = feat.size()
+    assert (len(size) == 4)
+    N, C = size[:2]
+    feat_var = feat.view(N, C, -1).var(dim=2) + eps
+    feat_std = feat_var.sqrt().view(N, C, 1, 1)
+    feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1)
+    return feat_mean, feat_std
+
+
+def adaptive_instance_normalization(content_feat, style_feat):
+    assert (content_feat.size()[:2] == style_feat.size()[:2])
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+
+    normalized_feat = (content_feat - content_mean.expand(
+        size)) / content_std.expand(size)
+    return normalized_feat * style_std.expand(size) + style_mean.expand(size)
+
+## AdaMean
+def adaptive_mean_normalization(content_feat, style_feat):
+    assert (content_feat.size()[:2] == style_feat.size()[:2])
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+
+    normalized_feat = (content_feat - content_mean.expand(
+        size))
+    return normalized_feat + style_mean.expand(size)
+
+## AdaStd
+def adaptive_std_normalization(content_feat, style_feat):
+    assert (content_feat.size()[:2] == style_feat.size()[:2])
+    size = content_feat.size()
+    style_mean, style_std = calc_mean_std(style_feat)
+    content_mean, content_std = calc_mean_std(content_feat)
+
+    normalized_feat = (content_feat) / content_std.expand(size)
+    return normalized_feat * style_std.expand(size)
+
+## EFDM
+def exact_feature_distribution_matching(content_feat, style_feat):
+    assert (content_feat.size() == style_feat.size())
+    B, C, W, H = content_feat.size(0), content_feat.size(1), content_feat.size(2), content_feat.size(3)
+    value_content, index_content = torch.sort(content_feat.view(B,C,-1))  # sort conduct a deep copy here.
+    value_style, _ = torch.sort(style_feat.view(B,C,-1))  # sort conduct a deep copy here.
+    inverse_index = index_content.argsort(-1)
+    new_content = content_feat.view(B,C,-1) + (value_style.gather(-1, inverse_index) - content_feat.view(B,C,-1).detach())
+
+    return new_content.view(B, C, W, H)
+
+## HM
+def histogram_matching(content_feat, style_feat):
+    assert (content_feat.size() == style_feat.size())
+    B, C, W, H = content_feat.size(0), content_feat.size(1), content_feat.size(2), content_feat.size(3)
+    x_view = content_feat.view(-1, W,H)
+    image1_temp = match_histograms(np.array(x_view.detach().clone().cpu().float().transpose(0, 2)),
+                                   np.array(style_feat.view(-1, W, H).detach().clone().cpu().float().transpose(0, 2)),
+                                   multichannel=True)
+    image1_temp = torch.from_numpy(image1_temp).float().to(content_feat.device).transpose(0, 2).view(B, C, W, H)
+    return content_feat + (image1_temp - content_feat).detach()
+
+
+
+def _calc_feat_flatten_mean_std(feat):
+    # takes 3D feat (C, H, W), return mean and std of array within channels
+    assert (feat.size()[0] == 3)
+    assert (isinstance(feat, torch.FloatTensor))
+    feat_flatten = feat.view(3, -1)
+    mean = feat_flatten.mean(dim=-1, keepdim=True)
+    std = feat_flatten.std(dim=-1, keepdim=True)
+    return feat_flatten, mean, std
+
+
+def _mat_sqrt(x):
+    U, D, V = torch.svd(x)
+    return torch.mm(torch.mm(U, D.pow(0.5).diag()), V.t())
+
+
+def coral(source, target):
+    # assume both source and target are 3D array (C, H, W)
+    # Note: flatten -> f
+
+    source_f, source_f_mean, source_f_std = _calc_feat_flatten_mean_std(source)
+    source_f_norm = (source_f - source_f_mean.expand_as(
+        source_f)) / source_f_std.expand_as(source_f)
+    source_f_cov_eye = \
+        torch.mm(source_f_norm, source_f_norm.t()) + torch.eye(3)
+
+    target_f, target_f_mean, target_f_std = _calc_feat_flatten_mean_std(target)
+    target_f_norm = (target_f - target_f_mean.expand_as(
+        target_f)) / target_f_std.expand_as(target_f)
+    target_f_cov_eye = \
+        torch.mm(target_f_norm, target_f_norm.t()) + torch.eye(3)
+
+    source_f_norm_transfer = torch.mm(
+        _mat_sqrt(target_f_cov_eye),
+        torch.mm(torch.inverse(_mat_sqrt(source_f_cov_eye)),
+                 source_f_norm)
+    )
+
+    source_f_transfer = source_f_norm_transfer * \
+                        target_f_std.expand_as(source_f_norm) + \
+                        target_f_mean.expand_as(source_f_norm)
+
+    return source_f_transfer.view(source.size())

net.py

@@ -0,0 +1,198 @@
+import torch.nn as nn
+import torch
+from function import adaptive_mean_normalization as adamean
+from function import adaptive_std_normalization as adastd
+from function import adaptive_instance_normalization as adain
+from function import exact_feature_distribution_matching as efdm
+from function import histogram_matching as hm
+
+from function import calc_mean_std
+# import ipdb
+from skimage.exposure import match_histograms
+import numpy as np
+
+decoder = nn.Sequential(
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 256, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 128, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 128, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 64, (3, 3)),
+    nn.ReLU(),
+    nn.Upsample(scale_factor=2, mode='nearest'),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 64, (3, 3)),
+    nn.ReLU(),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 3, (3, 3)),
+)
+
+vgg = nn.Sequential(
+    nn.Conv2d(3, 3, (1, 1)),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(3, 64, (3, 3)),
+    nn.ReLU(),  # relu1-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 64, (3, 3)),
+    nn.ReLU(),  # relu1-2
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(64, 128, (3, 3)),
+    nn.ReLU(),  # relu2-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 128, (3, 3)),
+    nn.ReLU(),  # relu2-2
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(128, 256, (3, 3)),
+    nn.ReLU(),  # relu3-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 256, (3, 3)),
+    nn.ReLU(),  # relu3-4
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(256, 512, (3, 3)),
+    nn.ReLU(),  # relu4-1, this is the last layer used
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu4-4
+    nn.MaxPool2d((2, 2), (2, 2), (0, 0), ceil_mode=True),
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-1
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-2
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU(),  # relu5-3
+    nn.ReflectionPad2d((1, 1, 1, 1)),
+    nn.Conv2d(512, 512, (3, 3)),
+    nn.ReLU()  # relu5-4
+)
+
+
+class Net(nn.Module):
+    def __init__(self, encoder, decoder, style):
+        super(Net, self).__init__()
+        enc_layers = list(encoder.children())
+        self.enc_1 = nn.Sequential(*enc_layers[:4])  # input -> relu1_1
+        self.enc_2 = nn.Sequential(*enc_layers[4:11])  # relu1_1 -> relu2_1
+        self.enc_3 = nn.Sequential(*enc_layers[11:18])  # relu2_1 -> relu3_1
+        self.enc_4 = nn.Sequential(*enc_layers[18:31])  # relu3_1 -> relu4_1
+        self.decoder = decoder
+        self.mse_loss = nn.MSELoss()
+        self.style = style
+
+        # fix the encoder
+        for name in ['enc_1', 'enc_2', 'enc_3', 'enc_4']:
+            for param in getattr(self, name).parameters():
+                param.requires_grad = False
+
+    # extract relu1_1, relu2_1, relu3_1, relu4_1 from input image
+    def encode_with_intermediate(self, input):
+        results = [input]
+        for i in range(4):
+            func = getattr(self, 'enc_{:d}'.format(i + 1))
+            results.append(func(results[-1]))
+        return results[1:]
+
+    # extract relu4_1 from input image
+    def encode(self, input):
+        for i in range(4):
+            input = getattr(self, 'enc_{:d}'.format(i + 1))(input)
+        return input
+
+    def calc_content_loss(self, input, target):
+        assert (input.size() == target.size())
+        assert (target.requires_grad is False)
+        return self.mse_loss(input, target)
+
+    def calc_style_loss(self, input, target):
+        # ipdb.set_trace()
+        assert (input.size() == target.size())
+        assert (target.requires_grad is False)  ## first make sure which one require gradient and which one do not.
+        # print(input.requires_grad) ## True
+        input_mean, input_std = calc_mean_std(input)
+        target_mean, target_std = calc_mean_std(target)
+        if self.style == 'adain':
+            return self.mse_loss(input_mean, target_mean) + \
+                   self.mse_loss(input_std, target_std)
+        elif self.style == 'adamean':
+            return self.mse_loss(input_mean, target_mean)
+        elif self.style == 'adastd':
+            return self.mse_loss(input_std, target_std)
+        elif self.style == 'efdm':
+            B, C, W, H = input.size(0), input.size(1), input.size(2), input.size(3)
+            value_content, index_content = torch.sort(input.view(B, C, -1))
+            value_style, index_style = torch.sort(target.view(B, C, -1))
+            inverse_index = index_content.argsort(-1)
+            return self.mse_loss(input.view(B,C,-1), value_style.gather(-1, inverse_index))
+        elif self.style == 'hm':
+            B, C, W, H = input.size(0), input.size(1), input.size(2), input.size(3)
+            x_view = input.view(-1, W, H)
+            image1_temp = match_histograms(np.array(x_view.detach().clone().cpu().float().transpose(0, 2)),
+                                           np.array(target.view(-1, W, H).detach().clone().cpu().float().transpose(0,2)),
+                                           multichannel=True)
+            image1_temp = torch.from_numpy(image1_temp).float().to(input.device).transpose(0, 2).view(B, C, W, H)
+            return self.mse_loss(input.reshape(B, C, -1), image1_temp.reshape(B, C, -1))
+        else:
+            raise NotImplementedError
+
+    def forward(self, content, style, alpha=1.0):
+        assert 0 <= alpha <= 1
+        # ipdb.set_trace()
+        style_feats = self.encode_with_intermediate(style)
+        content_feat = self.encode(content)
+        # print(content_feat.requires_grad) False
+        # print(style_feats[-1].requires_grad) False
+        if self.style == 'adain':
+            t = adain(content_feat, style_feats[-1])
+        elif self.style == 'adamean':
+            t = adamean(content_feat, style_feats[-1])
+        elif self.style == 'adastd':
+            t = adastd(content_feat, style_feats[-1])
+        elif self.style == 'efdm':
+            t = efdm(content_feat, style_feats[-1])
+        elif self.style == 'hm':
+            t = hm(content_feat, style_feats[-1])
+        else:
+            raise NotImplementedError
+        t = alpha * t + (1 - alpha) * content_feat
+
+        g_t = self.decoder(t)
+        g_t_feats = self.encode_with_intermediate(g_t)
+
+        loss_c = self.calc_content_loss(g_t_feats[-1], t) ### final feature should be the same.
+        loss_s = self.calc_style_loss(g_t_feats[0], style_feats[0])
+        for i in range(1, 4):
+            loss_s += self.calc_style_loss(g_t_feats[i], style_feats[i])
+        return loss_c, loss_s

test.py

@@ -0,0 +1,252 @@
+import argparse
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from PIL import Image
+from torchvision import transforms
+from torchvision.utils import save_image
+import time
+import net
+from function import adaptive_instance_normalization, coral
+from function import adaptive_mean_normalization
+from function import adaptive_std_normalization
+from function import exact_feature_distribution_matching, histogram_matching
+
+def test_transform(size, crop):
+    transform_list = []
+    if size != 0:
+        transform_list.append(transforms.Resize(size))
+    if crop:
+        transform_list.append(transforms.CenterCrop(size))
+    transform_list.append(transforms.ToTensor())
+    transform = transforms.Compose(transform_list)
+    return transform
+
+
+def style_transfer(vgg, decoder, content, style, alpha=1.0,
+                   interpolation_weights=None, style_type='adain'):
+    assert (0.0 <= alpha <= 1.0)
+    content_f = vgg(content)
+    style_f = vgg(style)
+    if interpolation_weights:
+        _, C, H, W = content_f.size()
+        feat = torch.FloatTensor(1, C, H, W).zero_().to(device)
+        if style_type == 'adain':
+            base_feat = adaptive_instance_normalization(content_f, style_f)
+        elif style_type == 'adamean':
+            base_feat = adaptive_mean_normalization(content_f, style_f)
+        elif style_type == 'adastd':
+            base_feat = adaptive_std_normalization(content_f, style_f)
+        elif style_type == 'efdm':
+            base_feat = exact_feature_distribution_matching(content_f, style_f)
+        elif style_type == 'hm':
+            feat = histogram_matching(content_f, style_f)
+        else:
+            raise NotImplementedError
+        for i, w in enumerate(interpolation_weights):
+            feat = feat + w * base_feat[i:i + 1]
+        content_f = content_f[0:1]
+    else:
+        if style_type == 'adain':
+            feat = adaptive_instance_normalization(content_f, style_f)
+        elif style_type == 'adamean':
+            feat = adaptive_mean_normalization(content_f, style_f)
+        elif style_type == 'adastd':
+            feat = adaptive_std_normalization(content_f, style_f)
+        elif style_type == 'efdm':
+            feat = exact_feature_distribution_matching(content_f, style_f)
+        elif style_type == 'hm':
+            feat = histogram_matching(content_f, style_f)
+        else:
+            raise NotImplementedError
+    feat = feat * alpha + content_f * (1 - alpha)
+    return decoder(feat)
+
+
+parser = argparse.ArgumentParser()
+# Basic options
+parser.add_argument('--content', type=str,
+                    help='File path to the content image')
+parser.add_argument('--content_dir', type=str,
+                    help='Directory path to a batch of content images')
+parser.add_argument('--style', type=str,
+                    help='File path to the style image, or multiple style \
+                    images separated by commas if you want to do style \
+                    interpolation or spatial control')
+parser.add_argument('--style_dir', type=str,
+                    help='Directory path to a batch of style images')
+parser.add_argument('--vgg', type=str, default='pretrained/vgg_normalised.pth')
+parser.add_argument('--decoder', type=str, default='pretrained/efdm_decoder_iter_160000.pth.tar')
+parser.add_argument('--style_type', type=str, default='adain', help='adain | adamean | adastd | efdm')
+parser.add_argument('--test_style_type', type=str, default='', help='adain | adamean | adastd | efdm')
+# Additional options
+parser.add_argument('--content_size', type=int, default=512,
+                    help='New (minimum) size for the content image, \
+                    keeping the original size if set to 0')
+parser.add_argument('--style_size', type=int, default=512,
+                    help='New (minimum) size for the style image, \
+                    keeping the original size if set to 0')
+parser.add_argument('--crop', action='store_true',
+                    help='do center crop to create squared image')
+parser.add_argument('--save_ext', default='.jpg',
+                    help='The extension name of the output image')
+parser.add_argument('--output', type=str, default='output',
+                    help='Directory to save the output image(s)')
+parser.add_argument('--photo', action='store_true',
+                    help='apply on the photo style transfer')
+# Advanced options
+parser.add_argument('--preserve_color', action='store_true',
+                    help='If specified, preserve color of the content image')
+parser.add_argument('--alpha', type=float, default=1.0,
+                    help='The weight that controls the degree of \
+                             stylization. Should be between 0 and 1')
+parser.add_argument(
+    '--style_interpolation_weights', type=str, default='',
+    help='The weight for blending the style of multiple style images')
+
+args = parser.parse_args()
+if not args.test_style_type:
+    args.test_style_type = args.style_type
+
+print('Note: the style type: %s and the pre-trained model: %s should be consistent' % (args.style_type, args.decoder))
+print('The test style type is:', args.test_style_type)
+
+do_interpolation = False
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+output_dir = Path(args.output + '_' + args.style_type + '_' + args.test_style_type)
+output_dir.mkdir(exist_ok=True, parents=True)
+
+# Either --content or --contentDir should be given.
+assert (args.content or args.content_dir)
+if args.content:
+    content_paths = [Path(args.content)]
+else:
+    content_dir = Path(args.content_dir)
+    content_paths = [f for f in content_dir.glob('*')]
+
+# Either --style or --styleDir should be given.
+assert (args.style or args.style_dir)
+if args.style:
+    style_paths = args.style.split(',')
+    if len(style_paths) == 1:
+        style_paths = [Path(args.style)]
+    else:
+        do_interpolation = True
+        # assert (args.style_interpolation_weights != ''), \
+        #     'Please specify interpolation weights'
+        # weights = [int(i) for i in args.style_interpolation_weights.split(',')]
+        # interpolation_weights = [w / sum(weights) for w in weights]
+else:
+    style_dir = Path(args.style_dir)
+    style_paths = [f for f in style_dir.glob('*')]
+
+decoder = net.decoder
+vgg = net.vgg
+
+decoder.eval()
+vgg.eval()
+
+decoder.load_state_dict(torch.load(args.decoder))
+vgg.load_state_dict(torch.load(args.vgg))
+vgg = nn.Sequential(*list(vgg.children())[:31])
+
+vgg.to(device)
+decoder.to(device)
+
+content_tf = test_transform(args.content_size, args.crop)
+style_tf = test_transform(args.style_size, args.crop)
+
+timer = []
+for content_path in content_paths:
+    if do_interpolation:
+        # one content image, 4 style image
+        style = torch.stack([style_tf(Image.open(str(p))) for p in style_paths])
+        content = content_tf(Image.open(str(content_path))) \
+            .unsqueeze(0).expand_as(style)
+        style = style.to(device)
+        content = content.to(device)
+        list = []
+        steps = [1, 0.75, 0.5, 0.25, 0]
+        for i in steps:
+            for j in steps:
+                list.append([i*j, i*(1-j), (1-i)*j, (1-i)*(1-j)])
+        count = 1
+        for interpolation_weights in list:
+            with torch.no_grad():
+                output = style_transfer(vgg, decoder, content, style,
+                                        args.alpha, interpolation_weights, style_type=args.test_style_type)
+            output = output.cpu()
+            output_name = output_dir / '{:s}_interpolate_{:s}_{:s}'.format(
+                content_path.stem, str(count), args.save_ext)
+            save_image(output, str(output_name))
+            count+=1
+
+        #### content & style trade-off.
+        # alpha = [0.0, 0.25, 0.5, 0.75, 1.0]
+        # for style_path in style_paths:
+        #     content = content_tf(Image.open(str(content_path)))
+        #     style = style_tf(Image.open(str(style_path)))
+        #     if args.preserve_color:
+        #         style = coral(style, content)
+        #     style = style.to(device).unsqueeze(0)
+        #     content = content.to(device).unsqueeze(0)
+        #     ## replace the style image with Gaussian noise
+        #     # style.normal_(0,1)
+        #     # style = torch.rand(style.size()).to(device)
+        #     ### for paired images.
+        #     if args.photo:
+        #         if content_path.stem[2:] == style_path.stem[3:]:
+        #             for sample_alpha in alpha:
+        #                 with torch.no_grad():
+        #                     output = style_transfer(vgg, decoder, content, style,
+        #                                             sample_alpha, style_type=args.test_style_type)
+        #                 output = output.cpu()
+        #                 output_name = output_dir / '{:s}_stylized_{:s}{:s}{:s}'.format(
+        #                     content_path.stem, style_path.stem, str(sample_alpha), args.save_ext)
+        #                 save_image(output, str(output_name))
+        #     else:
+        #         for sample_alpha in alpha:
+        #             with torch.no_grad():
+        #                 output = style_transfer(vgg, decoder, content, style,
+        #                                         sample_alpha, style_type=args.test_style_type)
+        #             output = output.cpu()
+        #             output_name = output_dir / '{:s}_stylized_{:s}{:s}{:s}'.format(
+        #                 content_path.stem, style_path.stem, str(sample_alpha), args.save_ext)
+        #             save_image(output, str(output_name))
+    else:  # process one content and one style
+        for style_path in style_paths:
+            content = content_tf(Image.open(str(content_path)))
+            style = style_tf(Image.open(str(style_path)))
+            if args.preserve_color:
+                style = coral(style, content)
+            style = style.to(device).unsqueeze(0)
+            content = content.to(device).unsqueeze(0)
+            ## replace the style image with Gaussian noise
+            # style.normal_(0,1)
+            # style = torch.rand(style.size()).to(device)
+            ### for paired images.
+            if args.photo:
+                if content_path.stem[2:] == style_path.stem[3:]:
+                    with torch.no_grad():
+                        start_time = time.time()
+                        output = style_transfer(vgg, decoder, content, style,
+                                                args.alpha, style_type=args.test_style_type)
+                        timer.append(time.time() - start_time)
+                        print(timer)
+
+                    output = output.cpu()
+                    output_name = output_dir / '{:s}_stylized_{:s}{:s}'.format(
+                        content_path.stem, style_path.stem, args.save_ext)
+                    save_image(output, str(output_name))
+            else:
+                with torch.no_grad():
+                    output = style_transfer(vgg, decoder, content, style,
+                                            args.alpha, style_type=args.test_style_type)
+                output = output.cpu()
+                output_name = output_dir / '{:s}_stylized_{:s}{:s}'.format(
+                    content_path.stem, style_path.stem, args.save_ext)
+                save_image(output, str(output_name))
+print(torch.FloatTensor(timer).mean())