-
Notifications
You must be signed in to change notification settings - Fork 26
/
test.py
169 lines (145 loc) · 6.48 KB
/
test.py
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
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
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 net
from function import nor_mean_std, nor_mean
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, SCT, content, style, alpha=1.0,
interpolation_weights=None):
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)
base_feat = SCT(content_f, style_f)
for i, w in enumerate(interpolation_weights):
feat = feat + w * base_feat[i:i + 1]
content_f = content_f[0:1]
else:
feat = SCT(content_f, style_f)
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('--testing_mode', default='art',
help='Artistic or Photo-realistic')
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='models/vgg_normalised.pth')
parser.add_argument('--decoder', type=str, help='put the trained decoder here')
parser.add_argument('--SCT', type=str, help='put the trained SCT module here')
# 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)')
# 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()
do_interpolation = False
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
output_dir = Path(args.output)
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
network = net.Net(vgg, decoder, args.testing_mode)
SCT = network.SCT
SCT.eval()
decoder.eval()
vgg.eval()
decoder.load_state_dict(torch.load(args.decoder))
vgg.load_state_dict(torch.load(args.vgg))
SCT.load_state_dict(torch.load(args.SCT))
vgg = nn.Sequential(*list(vgg.children())[:31]) if args.testing_mode == 'art' else nn.Sequential(*list(vgg.children())[:18])
decoder = decoder if args.testing_mode == 'art' else nn.Sequential(*list(net.decoder.children())[10:])
vgg.to(device)
decoder.to(device)
SCT.to(device)
content_tf = test_transform(args.content_size, args.crop)
style_tf = test_transform(args.style_size, args.crop)
for content_path in content_paths:
if do_interpolation: # one content image, N 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)
with torch.no_grad():
output = style_transfer(vgg, decoder, SCT, content, style,
args.alpha, interpolation_weights)
output = output.cpu()
output_name = output_dir / '{:s}_interpolation{:s}'.format(
content_path.stem, 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)
with torch.no_grad():
output = style_transfer(vgg, decoder, SCT, content, style,
args.alpha)
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))