first commit

app.py

@@ -0,0 +1,60 @@
+import utils
+import torch 
+import gradio as gr
+import numpy as np
+from PIL import Image
+from network import ImageTransformNet_dpws
+from torch.autograd import Variable
+from torchvision import transforms
+
+with gr.Blocks() as demo:
+    with gr.Row():
+      gr.HTML('<h1 style="text-align: center;">스타일 변환기</h1>')
+
+    with gr.Row():
+        with gr.Column():
+            style_radio = gr.Radio(['La muse', 'Mosaic', 'Starry Night Crop', 'Wave Crop'], label='원하는 스타일 선택!')
+            image_input = gr.Image(label='콘텐츠 이미지')
+            convert_button = gr.Button('변환!')
+        
+        with gr.Column():
+            result_image = gr.Image(label='결과 이미지')
+
+    def transform_image(style, img):
+        dtype = torch.FloatTensor
+            
+        # content image
+        img_transform_512 = transforms.Compose([
+                # transforms.Scale(512),                  # scale shortest side to image_size
+                transforms.Resize(512),                  # scale shortest side to image_size
+                # transforms.CenterCrop(512),             # crop center image_size out
+                transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
+                utils.normalize_tensor_transform()      # normalize with ImageNet values
+        ])
+
+        content = Image.fromarray(img)
+        content = img_transform_512(content)
+        content = content.unsqueeze(0)
+        # content = Variable(content).type(dtype)
+        content = Variable(content.repeat(1, 1, 1, 1), requires_grad=False).type(dtype)
+
+        # load style model
+        model_folder_name = '_'.join(style.lower().split())
+        model_path = 'models/' + model_folder_name + '/compressed.model'
+        checkpoint_lw = torch.load(model_path)
+
+        style_model = ImageTransformNet_dpws().type(dtype)
+        style_model.load_state_dict((checkpoint_lw))
+
+        # process input image
+        stylized = style_model(content).cpu()
+        utils.save_image('results.jpg', stylized.data[0])
+        return 'results.jpg'  
+
+    convert_button.click(
+        transform_image,
+        inputs=[style_radio, image_input],
+        outputs=[result_image],
+    )
+
+demo.launch()

models/la_muse/compressed.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b4d90ccfbf9679128254df891aec7a2a460f34df5490ad370bb1d3e6a32cc17
+size 92080

models/mosaic/compressed.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:24ff3368b90ba4ce65e0efaf97876bca4148958dd506e84a568a9425142cd314
+size 92080

models/starry_night_crop/compressed.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e5a3665567b07638dcc713c1d8ea30c1148824ebce3eade677b07181787a03d
+size 92080

models/wave_crop/compressed.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3d0974a83e9d62be50d636e80c3be48a90893376c22da6dc1cd004697616eeed
+size 92080

network.py

@@ -0,0 +1,356 @@
+import torch
+import torch.nn as nn
+
+# Conv Layer
+class ConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, groups=1):
+        super(ConvLayer, self).__init__()
+        paddings = kernel_size // 2
+        self.reflection_pad = nn.ReflectionPad2d(paddings)
+        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride, groups=groups) #, padding) 
+        # self.in_d = nn.InstanceNorm2d(out_channels, affine=True)
+
+    def forward(self, x):
+        out = self.reflection_pad(x)
+        out = self.conv2d(out)
+        return out
+
+
+class ConvLayer_dpws(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super(ConvLayer_dpws, self).__init__()
+        self.conv1 = ConvLayer(in_channels, in_channels, kernel_size, stride=stride, groups=in_channels)
+        self.in_1d = nn.InstanceNorm2d(in_channels, affine=True)
+        self.conv2 = ConvLayer(in_channels, out_channels, kernel_size=1, stride=1)
+        self.in_2d = nn.InstanceNorm2d(out_channels, affine=True)
+        self.relu = nn.ReLU()
+
+
+    def forward(self, x):
+        out = self.in_1d(self.conv1(x))
+        out = self.relu(self.in_2d(self.conv2(out)))
+        return out
+
+
+class ConvLayer_dpws_last(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super(ConvLayer_dpws_last, self).__init__()
+        self.conv1 = ConvLayer(in_channels, in_channels, kernel_size, stride=stride, groups=in_channels)
+        self.in_1d = nn.InstanceNorm2d(in_channels, affine=True)
+        self.conv2 = ConvLayer(in_channels, out_channels, kernel_size=1, stride=1)
+        # self.in_2d = nn.InstanceNorm2d(out_channels, affine=True)
+        # self.relu = nn.ReLU()
+
+    def forward(self, x):
+        out = self.in_1d(self.conv1(x))
+        # out = self.relu(self.in_2d(self.conv2(out)))
+        out = self.conv2(out)
+        return out
+
+# Upsample Conv Layer
+class UpsampleConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):
+        super(UpsampleConvLayer, self).__init__()
+        self.upsample = upsample
+        if upsample:
+            self.upsample = nn.Upsample(scale_factor=upsample, mode='nearest')
+        reflection_padding = kernel_size // 2
+        self.reflection_pad = nn.ReflectionPad2d(reflection_padding)
+        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride)
+        # self.in_d = nn.InstanceNorm2d(out_channels, affine=True)
+
+    def forward(self, x):
+        if self.upsample:
+            x = self.upsample(x)
+        out = self.reflection_pad(x)
+        out = self.conv2d(out)
+        return out
+
+
+class UpsampleConvLayer_dpws(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride, upsample=None):
+        super(UpsampleConvLayer_dpws, self).__init__()
+        self.upsample = upsample
+        if upsample:
+            self.upsample = nn.Upsample(scale_factor=upsample, mode='nearest')
+        self.conv1 = ConvLayer(in_channels, in_channels, kernel_size, stride, groups=in_channels)
+        self.in1 = nn.InstanceNorm2d(in_channels, affine=True )
+        self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1)
+        self.in2 = nn.InstanceNorm2d(out_channels, affine=True )
+        self.relu = nn.ReLU()
+
+    def forward(self, x):
+        if self.upsample:
+            x = self.upsample(x)
+        # out = self.reflection_pad(x)
+        # out = self.conv2d(out)
+        out = self.relu(self.in1(self.conv1(x)))
+        out = self.in2(self.conv2(out))
+        return out
+
+
+class DeConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super(DeConvLayer, self).__init__()
+
+        # reflection_padding = kernel_size // 2
+        # self.reflection_pad = nn.ReflectionPad2d(reflection_padding)
+        self.deconv2d = nn.ConvTranspose2d(in_channels, out_channels, kernel_size, stride, padding=1, output_padding=1)
+    def forward(self, x):
+        # out = self.reflection_pad(x)
+        out = self.deconv2d(x)
+        return out
+
+# Residual Block
+#   adapted from pytorch tutorial
+#   https://github.com/yunjey/pytorch-tutorial/blob/master/tutorials/02-
+#   intermediate/deep_residual_network/main.py
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+       
+        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1)
+        self.in1 = nn.InstanceNorm2d(channels, affine=True)
+        self.relu = nn.ReLU()
+        self.conv2 = ConvLayer(channels, channels, kernel_size=3, stride=1)
+        self.in2 = nn.InstanceNorm2d(channels, affine=True)
+
+
+    def forward(self, x):
+        residual = x
+        out = self.relu(self.in1(self.conv1(x)))
+        # out = self.relu(self.in2(self.conv2(out)))
+        out = self.in2(self.conv2(out))
+        # out = self.relu(self.conv2(out))
+        # out = self.conv2(out)
+        out = out + residual
+        # out = self.relu(out)
+       
+        return out 
+
+
+class ResidualBlock_depthwise(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock_depthwise, self).__init__()
+       
+        # ########################## deptwise ###########################################
+        self.conv1 = ConvLayer(channels, channels, kernel_size=3, stride=1, groups=channels)
+        self.in1 = nn.InstanceNorm2d(channels, affine=True )
+        self.conv2 = nn.Conv2d(channels, channels, kernel_size=1, stride=1)
+        self.in2 = nn.InstanceNorm2d(channels, affine=True )
+
+        self.conv3 = ConvLayer(channels, channels, kernel_size=3, stride=1, groups=channels)
+        self.in3 = nn.InstanceNorm2d(channels, affine=True )
+        self.conv4 = nn.Conv2d(channels, channels, kernel_size=1, stride=1)
+        self.in4 = nn.InstanceNorm2d(channels, affine=True )
+
+        self.relu = nn.ReLU()
+        self.prelu = nn.PReLU()
+
+    def forward(self, x):
+        
+        # ############### DEPTWISE ###################
+        # residual = x
+        # out = self.relu(self.in1(self.conv1(x)))
+        # out = self.relu(self.in2(self.conv2(out)))
+        # out = self.relu(self.in3(self.conv3(out)))
+        # out = self.relu(self.in4(self.conv4(out)))
+        # out = out + residual
+         
+        # # ################## v1 ####################
+        # residual = x
+        # out = self.in1(self.conv1(x))
+        # out = self.relu(self.in2(self.conv2(out)))
+        # out = self.in3(self.conv3(out))
+        # out = self.in4(self.conv4(out))
+        # out = out + residual
+        # out = self.relu(out) 
+
+        # ################## v2 ####################  √
+        residual = x
+        out = self.in1(self.conv1(x))
+        out = self.relu(self.in2(self.conv2(out)))
+        out = self.in3(self.conv3(out))
+        out = self.in4(self.conv4(out))
+        out = out + residual 
+
+        # ################## v3 ####################
+        # residual = x
+        # out = self.conv1(x)
+        # out = self.relu(self.in2(self.conv2(out)))
+        # out = self.conv3(out)
+        # out = self.in4(self.conv4(out))
+        # out = out + residual
+
+        # ################## v4 ####################
+        # residual = x
+        # out = self.in1(self.conv1(x))
+        # out = self.relu(self.in2(self.conv2(out)))
+        # out = self.in3(self.conv3(out))
+        # out = self.relu(self.in4(self.conv4(out)))
+        # out = out + residual
+        
+        return out 
+
+
+# Image Transform Network
+class ImageTransformNet(nn.Module):
+    def __init__(self):
+        super(ImageTransformNet, self).__init__()
+        
+        # nonlineraity
+        self.relu = nn.ReLU()
+        self.tanh = nn.Tanh()
+
+        # encoding layers
+        self.conv1 = ConvLayer(3, 32, kernel_size=9, stride=1)
+        self.in1_e = nn.InstanceNorm2d(32, affine=True )
+
+        self.conv2 = ConvLayer(32, 64, kernel_size=3, stride=2)
+        self.in2_e = nn.InstanceNorm2d(64, affine=True )
+
+        self.conv3 = ConvLayer(64, 128, kernel_size=3, stride=2)
+        self.in3_e = nn.InstanceNorm2d(128, affine=True )
+
+        # residual layers
+        self.res1 = ResidualBlock(128)
+        self.res2 = ResidualBlock(128)
+        self.res3 = ResidualBlock(128)
+        self.res4 = ResidualBlock(128)
+        self.res5 = ResidualBlock(128)
+        # self.res6 = ResidualBlock(128)
+
+        # decoding layers
+        # TODO:
+        # self.deconv3 = DeConvLayer(128, 64, kernel_size=3, stride=2)
+        self.deconv3 = UpsampleConvLayer(128, 64, kernel_size=3, stride=1, upsample=2)
+        self.in3_d = nn.InstanceNorm2d(64, affine=True )
+
+        # self.deconv2 = DeConvLayer(64, 32, kernel_size=3, stride=2)
+        self.deconv2 = UpsampleConvLayer(64, 32, kernel_size=3, stride=1, upsample=2)
+        self.in2_d = nn.InstanceNorm2d(32, affine=True )
+
+        self.deconv1 = ConvLayer(32, 3, kernel_size=9, stride=1)
+        self.in1_d = nn.InstanceNorm2d(3, affine=True )
+
+    def forward(self, x):
+        # encode
+        y = self.relu(self.in1_e(self.conv1(x)))
+        y = self.relu(self.in2_e(self.conv2(y)))
+        y = self.relu(self.in3_e(self.conv3(y)))
+        # y = self.relu(self.conv1(x))
+        # y = self.relu(self.conv2(y))
+        # y = self.relu(self.conv3(y))
+        y_downsample = y
+
+        # residual layers
+        y = self.res1(y)
+        y = self.res2(y)
+        y = self.res3(y)
+        y = self.res4(y)
+        y = self.res5(y)
+        
+        y_upsample = y
+        # decode
+        y = self.relu(self.in3_d(self.deconv3(y)))
+        y = self.relu(self.in2_d(self.deconv2(y)))
+        # y = self.relu(self.deconv3(y))
+        # y = self.relu(self.deconv2(y))
+        # y = self.tanh(self.in1_d(self.deconv1(y)))
+        y = self.deconv1(y)
+
+        # return y, y_downsample, y_upsample
+        return y
+
+
+ALAPHA_1 = 0.25
+ALAPHA_2 = 0.25
+# ALAPHA_1 = 0.5
+# ALAPHA_2 = 0.5
+class ImageTransformNet_dpws(nn.Module):
+    def __init__(self):
+        super(ImageTransformNet_dpws, self).__init__()
+        
+        # nonlineraity
+        self.relu = nn.ReLU()
+        self.tanh = nn.Tanh()
+
+        # encoding layers
+        self.conv1 = ConvLayer_dpws(3, int(32*ALAPHA_1), kernel_size=9, stride=1)
+        # self.in1_e = nn.InstanceNorm2d(int(32*ALAPHA_1), affine=True )
+        
+        self.conv2 = ConvLayer_dpws(int(32*ALAPHA_1), int(64*ALAPHA_1), kernel_size=3, stride= 2)
+        self.conv3 = ConvLayer_dpws(int(64*ALAPHA_1), int(128*ALAPHA_2), kernel_size=3, stride= 2)
+
+        # residual layers
+        self.res1 = ResidualBlock_depthwise(int(128*ALAPHA_2))
+        self.res2 = ResidualBlock_depthwise(int(128*ALAPHA_2))
+        self.res3 = ResidualBlock_depthwise(int(128*ALAPHA_2))
+        self.res4 = ResidualBlock_depthwise(int(128*ALAPHA_2))
+        self.res5 = ResidualBlock_depthwise(int(128*ALAPHA_2))
+        # self.res6 = ResidualBlock_depthwise(128)
+
+        # decoding layers
+        # TODO:
+        # self.deconv3 = DeConvLayer(128, 64, kernel_size=3, stride=2)
+        self.deconv3 = UpsampleConvLayer_dpws(int(128*ALAPHA_2), int(64*ALAPHA_1), kernel_size=3, stride=1, upsample=2)
+        # self.in3_d = nn.InstanceNorm2d(int(64*ALAPHA_1), affine=True )
+
+        # self.deconv2 = DeConvLayer(64, 32, kernel_size=3, stride=2)
+        self.deconv2 = UpsampleConvLayer_dpws(int(64*ALAPHA_1), int(32*ALAPHA_1), kernel_size=3, stride=1, upsample=2)
+        # self.in2_d = nn.InstanceNorm2d(32, affine=True )
+
+        self.deconv1 = ConvLayer_dpws_last(int(32*ALAPHA_1), 3, kernel_size=9, stride=1)
+        # self.deconv1 = ConvLayer_dpws_last(int(32*ALAPHA_1), 3, kernel_size=9, stride=1)
+        # self.in1_d = nn.InstanceNorm2d(3, affine=True )
+
+    def forward(self, x):
+        # encode
+        # y = self.relu(self.in1_e(self.conv1(x)))
+        y = self.conv1(x)
+        y = self.conv2(y)
+        y = self.conv3(y)
+        y_downsample = y
+
+        # y = self.relu(self.in2_e(self.conv2(y)))
+        # y = self.relu(self.in3_e(self.conv3(y)))
+        # residual layers
+        y = self.res1(y)
+        y = self.res2(y)
+        y = self.res3(y)
+        y = self.res4(y)
+        y = self.res5(y)
+        # y = self.res6(y)
+        y_upsample = y
+
+        # decode
+        y = self.deconv3(y)
+        y = self.deconv2(y)
+        y = self.deconv1(y)
+
+        # return y, y_downsample, y_upsample
+        return y
+
+class distiller_1(nn.Module):
+    def __init__(self):
+        super(distiller_1, self).__init__()
+        
+        self.conv = nn.Conv2d(128, int(128*ALAPHA_2), kernel_size=1, stride=1)
+
+    def forward(self, x):
+        # encode
+        y = self.conv(x)
+
+        return y
+
+class distiller_2(nn.Module):
+    def __init__(self):
+        super(distiller_2, self).__init__()
+        
+        self.conv = nn.Conv2d(128, int(128*ALAPHA_2), kernel_size=1, stride=1)
+
+    def forward(self, x):
+        # encode
+        y = self.conv(x)
+
+        return y

requirements.txt

@@ -0,0 +1,5 @@
+torch==1.12.1
+torchvision==0.13.1
+numpy==1.24.0
+torchsummary
+ptflops

style.py

@@ -0,0 +1,258 @@
+import numpy as np
+import torch
+import os
+import argparse
+import time
+import collections
+
+from torch.autograd import Variable
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision import transforms
+
+import utils
+from network import ImageTransformNet, ImageTransformNet_dpws
+from vgg import Vgg16
+
+# Global Variables
+IMAGE_SIZE = 256
+BATCH_SIZE = 4
+LEARNING_RATE = 1e-3
+EPOCHS = 2
+
+# STYLE_WEIGHT = 9.2e3
+# STYLE_WEIGHT = 8e4
+STYLE_WEIGHT = 7e4
+# STYLE_WEIGHT = 0
+# CONTENT_WEIGHT = 1e-2
+# CONTENT_WEIGHT = 0.15
+CONTENT_WEIGHT = 0.1
+L1_WEIGHT = 1
+
+def train(args):
+    # GPU enabling
+    if (args.gpu != None):
+        use_cuda = True
+        dtype = torch.cuda.FloatTensor
+        torch.cuda.set_device(args.gpu)
+        print("Current device: %d" %torch.cuda.current_device())
+
+    # visualization of training controlled by flag
+    visualize = (args.visualize != None)
+    if (visualize):
+        img_transform_512 = transforms.Compose([
+            transforms.Scale(512),                  # scale shortest side to image_size
+            # transforms.CenterCrop(512),           # crop center image_size out
+            transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
+            utils.normalize_tensor_transform()      # normalize with ImageNet values
+        ])
+
+
+        testImage_maine = utils.load_image(args.test_image)
+        testImage_maine = img_transform_512(testImage_maine)
+        testImage_maine = Variable(testImage_maine.repeat(1, 1, 1, 1), requires_grad=False).type(dtype)
+        test_name = os.path.split(args.test_image)[-1].split('.')[0]
+
+    # define network
+    image_transformer_dpws = ImageTransformNet_dpws().type(dtype)
+    # paras = [image_transformer_dpws.parameters()]
+    optimizer = Adam(image_transformer_dpws.parameters(), LEARNING_RATE) 
+
+    loss_mse = torch.nn.MSELoss()
+    loss_l1 = torch.nn.L1Loss() 
+   
+    vgg = Vgg16().type(dtype)
+    image_transformer = ImageTransformNet().type(dtype)
+    image_transformer.load_state_dict(torch.load(args.load_path))
+
+    # get training dataset
+    dataset_transform = transforms.Compose([
+        transforms.Scale(IMAGE_SIZE),           # scale shortest side to image_size
+        transforms.CenterCrop(IMAGE_SIZE),      # crop center image_size out
+        transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
+        utils.normalize_tensor_transform()      # normalize with ImageNet values
+    ])
+    train_dataset = datasets.ImageFolder(args.dataset, dataset_transform)
+    train_loader = DataLoader(train_dataset, batch_size = BATCH_SIZE)
+
+    # style image
+    style_transform = transforms.Compose([
+        transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
+        utils.normalize_tensor_transform()      # normalize with ImageNet values
+    ])
+    style = utils.load_image(args.style_image)
+    style = style_transform(style)
+    style = Variable(style.repeat(BATCH_SIZE, 1, 1, 1)).type(dtype)
+    style_name = os.path.split(args.style_image)[-1].split('.')[0]
+    
+    # calculate gram matrices for style feature layer maps we care about
+    style_features = vgg(style)
+    style_gram = [utils.gram(fmap) for fmap in style_features]
+
+    for e in range(EPOCHS):
+
+        # track values for...
+        img_count = 0
+        aggregate_style_loss = 0.0
+        aggregate_content_loss = 0.0
+        aggregate_l1_loss = 0.0
+        # aggregate_tv_loss = 0.0
+
+        # train network
+        image_transformer_dpws.train()
+        for batch_num, (x, label) in enumerate(train_loader):
+            img_batch_read = len(x)
+            img_count += img_batch_read
+
+            # zero out gradients
+            optimizer.zero_grad()
+
+            # input batch to transformer network
+
+            x = Variable(x).type(dtype)
+            y_hat = image_transformer_dpws(x)
+            y_label = image_transformer(x)
+
+            # get vgg features
+            y_c_features = vgg(x)
+            y_hat_features = vgg(y_hat)
+
+            # calculate style loss
+            y_hat_gram = [utils.gram(fmap) for fmap in y_hat_features]
+            style_loss = 0.0
+            
+            for j in range(4):
+                style_loss += loss_mse(y_hat_gram[j], style_gram[j][:img_batch_read])
+            style_loss = STYLE_WEIGHT*style_loss
+            aggregate_style_loss += style_loss.data.item()
+
+            # calculate content loss (h_relu_2_2)
+            recon = y_c_features[1]      
+            recon_hat = y_hat_features[1]
+            content_loss = CONTENT_WEIGHT*loss_mse(recon_hat, recon)
+            aggregate_content_loss += content_loss.data.item()
+
+            # calculate l1 loss
+            l1_loss = L1_WEIGHT*loss_mse(y_hat, y_label)
+            aggregate_l1_loss += l1_loss.data.item()
+
+            # total loss
+            # total_loss = style_loss + content_loss + tv_loss + l1_loss + dis_loss
+            total_loss = style_loss + l1_loss + content_loss
+
+            # backprop
+            total_loss.backward()
+            optimizer.step()
+
+            # print out status message
+            if ((batch_num + 1) % 100 == 0):
+                status = "{}  Epoch {}:  [{}/{}]  Batch:[{}]  agg_style: {:.6f}  agg_l1: {:.6f}  agg_content: {:.6f} ".format(
+                                time.ctime(), e + 1, img_count, len(train_dataset), batch_num+1,
+                                aggregate_style_loss/(batch_num+1.0), aggregate_l1_loss/(batch_num+1.0), aggregate_content_loss/(batch_num+1.0)
+                                )
+                print(status)
+
+            if ((batch_num + 1) % 5000 == 0) and (visualize):
+                image_transformer_dpws.eval()
+
+                if not os.path.exists("visualization"):
+                    os.makedirs("visualization")
+
+                outputTestImage_maine = image_transformer_dpws(testImage_maine)
+    
+                test_path = "visualization/%s/%s%d_%05d.jpg" %(style_name, test_name, e+1, batch_num+1)
+                utils.save_image(test_path, outputTestImage_maine.data[0].cpu())
+
+                print("images saved")
+                image_transformer_dpws.train()
+
+    # save model
+    image_transformer_dpws.eval()
+
+    if use_cuda:
+        image_transformer_dpws.cpu()
+
+    if not os.path.exists("models"):
+        os.makedirs("models")
+    filename = "models/%s.model" %style_name
+    torch.save(image_transformer_dpws.state_dict(), filename)
+    
+    if use_cuda:
+        image_transformer_dpws.cuda()
+
+def style_transfer(args):
+    # GPU enabling
+    if (args.gpu != None):
+        use_cuda = True
+        dtype = torch.cuda.FloatTensor
+        torch.cuda.set_device(args.gpu)
+        print("Current device: %d" %torch.cuda.current_device())
+    else : 
+        dtype = torch.FloatTensor
+        
+    # content image
+    img_transform_512 = transforms.Compose([
+            # transforms.Scale(512),                  # scale shortest side to image_size
+            transforms.Resize(512),                  # scale shortest side to image_size
+            # transforms.CenterCrop(512),             # crop center image_size out
+            transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
+            utils.normalize_tensor_transform()      # normalize with ImageNet values
+    ])
+
+    content = utils.load_image(args.source)
+    content = img_transform_512(content)
+    content = content.unsqueeze(0)
+    # content = Variable(content).type(dtype)
+    content = Variable(content.repeat(1, 1, 1, 1), requires_grad=False).type(dtype)
+
+    # load style model
+    checkpoint_lw = torch.load(args.model_path)
+
+    style_model = ImageTransformNet_dpws().type(dtype)
+    style_model.load_state_dict((checkpoint_lw))
+
+    # process input image
+    stylized = style_model(content).cpu()
+    utils.save_image(args.output, stylized.data[0])
+
+
+def main():
+    parser = argparse.ArgumentParser(description='style transfer in pytorch')
+    subparsers = parser.add_subparsers(title="subcommands", dest="subcommand")
+
+    train_parser = subparsers.add_parser("train", help="train a model to do style transfer")
+    train_parser.add_argument("--style_image", type=str, required=True, help="path to a style image to train with")
+    train_parser.add_argument("--test_image", type=str, required=True, help="path to a test image to test with")
+    train_parser.add_argument("--dataset", type=str, required=True, help="path to a dataset")
+    train_parser.add_argument("--gpu", type=int, default=None, help="ID of GPU to be used")
+    train_parser.add_argument("--visualize", type=int, default=None, help="Set to 1 if you want to visualize training")
+
+    style_parser = subparsers.add_parser("transfer", help="do style transfer with a trained model")
+    style_parser.add_argument("--model_path", type=str, required=True, help="path to a pretrained model for a style image")
+    style_parser.add_argument("--source", type=str, required=True, help="path to source image")
+    style_parser.add_argument("--output", type=str, required=True, help="file name for stylized output image")
+    style_parser.add_argument("--gpu", type=int, default=None, help="ID of GPU to be used")
+
+    args = parser.parse_args()
+
+    # command
+    if (args.subcommand == "train"):
+        print("Training!")
+        train(args)
+    elif (args.subcommand == "transfer"):
+        print("Style transfering!")
+        style_transfer(args)
+    else:
+        print("invalid command")
+
+if __name__ == '__main__':
+    main()
+
+
+
+
+
+
+
+

utils.py

@@ -0,0 +1,32 @@
+import torch
+from PIL import Image
+from torch.autograd import Variable
+from torchvision import transforms
+import numpy as np
+
+# opens and returns image file as a PIL image (0-255)
+def load_image(filename):
+    img = Image.open(filename)
+    return img
+
+# assumes data comes in batch form (ch, h, w)
+def save_image(filename, data):
+    std = np.array([0.229, 0.224, 0.225]).reshape((3, 1, 1))
+    mean = np.array([0.485, 0.456, 0.406]).reshape((3, 1, 1))
+    img = data.clone().numpy()
+    img = ((img * std + mean).transpose(1, 2, 0)*255.0).clip(0, 255).astype("uint8")
+    img = Image.fromarray(img)
+    img.save(filename)
+
+# Calculate Gram matrix (G = FF^T)
+def gram(x):
+    (bs, ch, h, w) = x.size()
+    f = x.view(bs, ch, w*h)
+    f_T = f.transpose(1, 2)
+    G = f.bmm(f_T) / (ch * h * w)
+    return G
+
+# using ImageNet values
+def normalize_tensor_transform():
+    return transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])