Skip to content

Commit cfc84b1

Browse files
wang-zywang-zy
committed
Finish naive and realtime ST
1 parent 837e4e2 commit cfc84b1

11 files changed

+523
-1089
lines changed

Diff for: ‎CompareModels.py

+122
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
import torch
2+
import torchvision.models as models
3+
from torch.autograd import Variable
4+
import numpy as np
5+
import torch.nn as nn
6+
import torch.nn.functional as F
7+
from PIL import Image
8+
from torchvision import transforms
9+
import glob
10+
from torch.utils.data.dataset import Dataset
11+
12+
13+
def normalize(x):
14+
y = x.div(255)
15+
mean = [0.485, 0.456, 0.406]
16+
std = [0.229, 0.224, 0.225]
17+
y[:, 0, :, :] = (y[:, 0, :, :] - mean[0]) / std[0]
18+
y[:, 1, :, :] = (y[:, 1, :, :] - mean[1]) / std[1]
19+
y[:, 2, :, :] = (y[:, 2, :, :] - mean[2]) / std[2]
20+
return y
21+
22+
23+
class Conv(nn.Module):
24+
def __init__(self, in_channels, out_channels, kernel_size, stride=1):
25+
super(Conv, self).__init__()
26+
reflection_padding = kernel_size // 2
27+
self.reflection_pad = nn.ReflectionPad2d(reflection_padding)
28+
self.conv2d = nn.Conv2d(in_channels, out_channels, kernel_size, stride=stride)
29+
30+
def forward(self, x):
31+
return self.conv2d(self.reflection_pad(x))
32+
33+
34+
class Res(nn.Module):
35+
def __init__(self, numChannels):
36+
super(Res, self).__init__()
37+
bn_flag = True
38+
self.conv1 = Conv(numChannels, numChannels, 3, stride=1)
39+
self.in1 = nn.InstanceNorm2d(numChannels, affine=bn_flag)
40+
self.relu = nn.ReLU()
41+
self.conv2 = Conv(numChannels, numChannels, 3, stride=1)
42+
self.in2 = nn.InstanceNorm2d(numChannels, affine=bn_flag)
43+
def forward(self, x):
44+
residual = x
45+
output = self.in2(self.conv2(self.relu(self.in1(self.conv1(x)))))
46+
return residual + output
47+
48+
49+
class DeConv(nn.Module):
50+
def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, upsample=2):
51+
super(DeConv, self).__init__()
52+
self.upsample = nn.Upsample(scale_factor=upsample)
53+
self.conv = Conv(in_channels, out_channels, kernel_size, stride=stride)
54+
55+
def forward(self, x):
56+
return self.conv(self.upsample(x))
57+
58+
59+
class StyleNet(nn.Module):
60+
def __init__(self):
61+
super(StyleNet, self).__init__()
62+
bn_flag = True
63+
self.relu = nn.ReLU()
64+
self.conv1 = Conv(3, 32, 9, stride=1)
65+
self.in1 = nn.InstanceNorm2d(32, affine=bn_flag)
66+
self.conv2 = Conv(32, 64, 3, stride=2)
67+
self.in2 = nn.InstanceNorm2d(64, affine=bn_flag)
68+
self.conv3 = Conv(64, 128, 3, stride=2)
69+
self.in3 = nn.InstanceNorm2d(128, affine=bn_flag)
70+
self.res1 = Res(128)
71+
self.res2 = Res(128)
72+
self.res3 = Res(128)
73+
self.res4 = Res(128)
74+
self.res5 = Res(128)
75+
self.deconv1 = DeConv(128, 64, kernel_size=3, stride=1, upsample=2)
76+
self.in4 = nn.InstanceNorm2d(64, affine=bn_flag)
77+
self.deconv2 = DeConv(64, 32, kernel_size=3, stride=1, upsample=2)
78+
self.in5 = nn.InstanceNorm2d(32, affine=bn_flag)
79+
self.conv4 = Conv(32, 3, 9, stride=1)
80+
81+
def forward(self, x):
82+
x = self.relu(self.in1(self.conv1(x)))
83+
x = self.relu(self.in2(self.conv2(x)))
84+
x = self.relu(self.in3(self.conv3(x)))
85+
x = self.res5(self.res4(self.res3(self.res2(self.res1(x)))))
86+
x = self.relu(self.in4(self.deconv1(x)))
87+
x = self.relu(self.in5(self.deconv2(x)))
88+
return self.conv4(x)
89+
90+
def LoadImage(fname, scale=False):
91+
# load image and convert to tensor wrapped in a variable
92+
if scale is True:
93+
loader = transforms.Compose([transforms.Scale((imsize, imsize)),
94+
transforms.CenterCrop(imsize),
95+
transforms.ToTensor(),
96+
transforms.Lambda(lambda x: x.mul(255))])
97+
else:
98+
loader = transforms.Compose([transforms.Scale((imsize, imsize)),
99+
transforms.CenterCrop(imsize),
100+
transforms.ToTensor()])
101+
image = Image.open(fname).convert('RGB')
102+
data = loader(image)
103+
data = Variable(data.cuda(), volatile=True)
104+
data = data.unsqueeze(0)
105+
return data
106+
107+
108+
def SaveImage(tensor_transformed, fname, factor=255):
109+
def RGB(image):
110+
return (image.transpose(0, 2, 3, 1)*factor).clip(0, 255).astype(np.uint8)
111+
result = Image.fromarray(RGB(tensor_transformed.data.cpu().numpy())[0])
112+
result.save(fname)
113+
114+
115+
imsize = 256
116+
117+
gen = glob.glob('SavedModels/*.model')
118+
s = torch.load(gen[0])
119+
print(gen[0])
120+
image = LoadImage('amber.jpg', scale=True)
121+
SaveImage(s(image), 'SavedImages/candy2.png', factor=1)
122+

Diff for: ‎DataPreProc.py

+8-84
Original file line numberDiff line numberDiff line change
@@ -6,93 +6,17 @@
66
import torch
77
from os import listdir
88

9-
##############################################
10-
#
11-
# Preprecess the COCO dataset
12-
# (run different section separately)
13-
# 1. Resize to 256x256
14-
# 2. Remove all grey scale images
15-
# 3. Calculate mean values for each channel
16-
# 4. Calculate std for each channel
17-
#
18-
##############################################
199

20-
21-
# # resize COCO dataset
22-
# dir = 'Dataset/train2014/'
23-
# gen = listdir(dir)
24-
25-
# for i in range(len(gen)):
26-
# fname = dir + gen[i]
27-
# im = Image.open(fname)
28-
# img = im.resize((256, 256), Image.ANTIALIAS)
29-
# img.save('Dataset/PreData/'+str(i)+'.jpg', 'JPEG')
30-
31-
# get mean value of all samples
32-
gen = glob.glob('Dataset/PreData/*.jpg')
33-
34-
loader = transforms.ToTensor()
35-
36-
# # remove grey scale images
37-
# for i in range(len(gen)):
38-
# fname = gen[i]
39-
# im = Image.open(fname)
40-
# data = loader(im)
41-
#
42-
# if data.shape[0] != 3:
43-
# print i, gen[i]
44-
# os.remove(gen[i])
45-
46-
# r_mean = 0.0
47-
# g_mean = 0.0
48-
# b_mean = 0.0
49-
#
50-
# for i in range(len(gen)):
51-
# if i % 1000 == 0:
52-
# print i
53-
# fname = gen[i]
54-
# im = Image.open(fname)
55-
# data = loader(im)
56-
#
57-
# size = data.shape[1] * data.shape[2] * 1.0
58-
# r_mean += torch.sum(data[0, :, :]) / size
59-
# g_mean += torch.sum(data[1, :, :]) / size
60-
# b_mean += torch.sum(data[2, :, :]) / size
61-
#
62-
# r_mean /= len(gen)
63-
# g_mean /= len(gen)
64-
# b_mean /= len(gen)
65-
#
66-
# print r_mean
67-
# print g_mean
68-
# print b_mean
69-
70-
# calculate std for each dimension for all samples
71-
r_mean = 0.471116568055
72-
g_mean = 0.446522854439
73-
b_mean = 0.406828638561
74-
75-
r_var = 0.0
76-
g_var = 0.0
77-
b_var = 0.0
10+
# resize COCO dataset
11+
gen = glob.glob('Dataset/train2014/*.jpg')
12+
loader = transforms.Compose([transforms.Scale(256), transforms.CenterCrop(256), transforms.ToTensor()])
13+
f = transforms.ToPILImage()
7814

7915
for i in range(len(gen)):
8016
if i % 1000 == 0:
81-
print i
82-
im = Image.open(gen[i])
17+
print(i)
18+
im = Image.open(gen[i]).convert('RGB')
8319
data = loader(im)
84-
r_var += torch.var(data[0, :, :] - r_mean)
85-
g_var += torch.var(data[1, :, :] - g_mean)
86-
b_var += torch.var(data[2, :, :] - b_mean)
87-
88-
r_var /= len(gen)
89-
g_var /= len(gen)
90-
b_var /= len(gen)
91-
92-
print np.sqrt(r_var)
93-
print np.sqrt(g_var)
94-
print np.sqrt(b_var)
20+
img = f(data)
21+
img.save('Dataset/CropData/'+str(i)+'.jpg', 'JPEG')
9522

96-
r_std = 0.244088944037
97-
g_std = 0.239412672854
98-
b_std = 0.243587113539

0 commit comments

Comments
 (0)