Real-time-Image-Style-Transfer/CompareModels.py at master · wang-zy/Real-time-Image-Style-Transfer · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
import torch
import torchvision.models as models
from torch.autograd import Variable
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torchvision import transforms
import glob
from torch.utils.data.dataset import Dataset


def normalize(x):
    y = x.div(255)
    mean = [0.485, 0.456, 0.406]
    std = [0.229, 0.224, 0.225]
    y[:, 0, :, :] = (y[:, 0, :, :] - mean[0]) / std[0]
    y[:, 1, :, :] = (y[:, 1, :, :] - mean[1]) / std[1]
    y[:, 2, :, :] = (y[:, 2, :, :] - mean[2]) / std[2]
    return y


class Conv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1):
        super(Conv, self).__init__()
        reflection_padding = kernel_size // 2
        self.reflection_pad = nn.ReflectionPad2d(reflection_padding)
        self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride)

    def forward(self, x):
        return self.conv2d(self.reflection_pad(x))


class Res(nn.Module):
    def __init__(self, numChannels):
        super(Res, self).__init__()
        bn_flag = True
        self.conv1 = Conv(numChannels, numChannels, 3, stride=1)
        self.in1 = nn.InstanceNorm2d(numChannels, affine=bn_flag)
        self.relu = nn.ReLU()
        self.conv2 = Conv(numChannels, numChannels, 3, stride=1)
        self.in2 = nn.InstanceNorm2d(numChannels, affine=bn_flag)
    def forward(self, x):
        residual = x
        output = self.in2(self.conv2(self.relu(self.in1(self.conv1(x)))))
        return residual + output


class DeConv(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, upsample=2):
        super(DeConv, self).__init__()
        self.upsample = nn.Upsample(scale_factor=upsample)
        self.conv = Conv(in_channels, out_channels, kernel_size, stride=stride)

    def forward(self, x):
        return self.conv(self.upsample(x))


class StyleNet(nn.Module):
    def __init__(self):
        super(StyleNet, self).__init__()
        bn_flag = True
        self.relu = nn.ReLU()
        self.conv1 = Conv(3, 32, 9, stride=1)
        self.in1 = nn.InstanceNorm2d(32, affine=bn_flag)
        self.conv2 = Conv(32, 64, 3, stride=2)
        self.in2 = nn.InstanceNorm2d(64, affine=bn_flag)
        self.conv3 = Conv(64, 128, 3, stride=2)
        self.in3 = nn.InstanceNorm2d(128, affine=bn_flag)
        self.res1 = Res(128)
        self.res2 = Res(128)
        self.res3 = Res(128)
        self.res4 = Res(128)
        self.res5 = Res(128)
        self.deconv1 = DeConv(128, 64, kernel_size=3, stride=1, upsample=2)
        self.in4 = nn.InstanceNorm2d(64, affine=bn_flag)
        self.deconv2 = DeConv(64, 32, kernel_size=3, stride=1, upsample=2)
        self.in5 = nn.InstanceNorm2d(32, affine=bn_flag)
        self.conv4 = Conv(32, 3, 9, stride=1)

    def forward(self, x):
        x = self.relu(self.in1(self.conv1(x)))
        x = self.relu(self.in2(self.conv2(x)))
        x = self.relu(self.in3(self.conv3(x)))
        x = self.res5(self.res4(self.res3(self.res2(self.res1(x)))))
        x = self.relu(self.in4(self.deconv1(x)))
        x = self.relu(self.in5(self.deconv2(x)))
        return self.conv4(x)

def LoadImage(fname, scale=False):
    # load image and convert to tensor wrapped in a variable
    if scale is True:
        loader = transforms.Compose([transforms.Scale((imsize, imsize)),
                                     transforms.CenterCrop(imsize),
                                     transforms.ToTensor(),
                                     transforms.Lambda(lambda x: x.mul(255))])
    else:
        loader = transforms.Compose([transforms.Scale((imsize, imsize)),
                                     transforms.CenterCrop(imsize),
                                     transforms.ToTensor()])
    image = Image.open(fname).convert('RGB')
    data = loader(image)
    data = Variable(data.cuda(), volatile=True)
    data = data.unsqueeze(0)
    return data


def SaveImage(tensor_transformed, fname, factor=255):
    def RGB(image):
        return (image.transpose(0, 2, 3, 1)*factor).clip(0, 255).astype(np.uint8)
    result = Image.fromarray(RGB(tensor_transformed.data.cpu().numpy())[0])
    result.save(fname)


imsize = 256

gen = glob.glob('SavedModels/*.model')
s = torch.load(gen[0])
print(gen[0])
image = LoadImage('amber.jpg', scale=True)
SaveImage(s(image), 'SavedImages/candy2.png', factor=1)