def generate(model, args):
content_image = args['content_image'].convert('RGB')
style_image = args['style_image'].convert('RGB')
preserve_color = args['preserve_color']
alpha = args['alpha']
print('[GENERATE] Ran with preserve_color "{}". alpha "{}"'.format(preserve_color, alpha))
vgg = model['vgg']
decoder = model['decoder']
content_tf = model['content_tf']
style_tf = model['style_tf']
content = content_tf(content_image)
style = style_tf(style_image)
if 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, content, style, alpha)
ndarr = output[0].mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
return {
'image': Image.fromarray(ndarr)
}
def main_func(content_path, style_path):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
content_path = Path(content_path)
style_path = Path(style_path)
output_dir = Path('output')
output_dir.mkdir(exist_ok=True, parents=True)
do_interpolation = False
decoder = net.decoder
vgg = net.vgg
decoder.eval()
vgg.eval()
decoder.load_state_dict(torch.load('models/decoder.pth'))
vgg.load_state_dict(torch.load('models/vgg_normalised.pth'))
vgg = nn.Sequential(*list(vgg.children())[:31])
vgg.to(device)
decoder.to(device)
content_size = 1024
style_size = 500
crop = True
preserve_color = False
alpha = 1.0
interpolation_weights = [0.3, 0.7]
content_tf = test_transform(content_size, crop)
style_tf = test_transform(style_size, crop)
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, content, style, alpha,
interpolation_weights)
output = output.cpu()
else: # process one content and one style
content = content_tf(Image.open(str(content_path)))
style = style_tf(Image.open(str(style_path)))
if 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, content, style, alpha)
output = output.cpu()
output_name = output_dir / '{:s}_stylized_{:s}_alpha_{:s}{:s}'.format(
content_path.stem, style_path.stem, str(alpha), '.jpg')
save_image(output, str(output_name))
return output_name
def transfer(self,
content,
style,
preserve_color=False,
alpha=1.0,
interpolation_weights=None):
"""
CONTENT is always a single image.
STYLE can be either a single image or a list of images.
If STYLE is a list of images, you must also pass a list of INTERPOLATION_WEIGHTS.
"""
if interpolation_weights:
# one content image, N style images
style = torch.stack([self.style_tf(s) for s in style])
content = self.content_tf(content).unsqueeze(0).expand_as(style)
style = style.to(self.device)
content = content.to(self.device)
else:
# one content image, one style image
content = self.content_tf(content)
style = self.style_tf(style)
if preserve_color:
style = coral(style, content)
style = style.to(self.device).unsqueeze(0)
content = content.to(self.device).unsqueeze(0)
with torch.no_grad():
content_f = self.vgg(content)
style_f = self.vgg(style)
if interpolation_weights:
_, C, H, W = content_f.size()
feat = torch.FloatTensor(1, C, H, W).zero_().to(self.device)
base_feat = adaptive_instance_normalization(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 = adaptive_instance_normalization(content_f, style_f)
feat = feat * alpha + content_f * (1 - alpha)
output = self.decoder(feat)
return output.cpu()
def image_process(content, style):
content_tf1 = content_transform()
content_frame = content_tf1(content)
#content_frame = torch.tensor(content_frame)
h, w, c = np.shape(content_frame)
style_tf1 = style_transform(h, w)
style = style_tf1(style.convert("RGB"))
#style = torch.tensor(style)
if yaml['preserve_color']:
style = coral(style, content)
style = style.to(device).unsqueeze(0)
content = content_frame.to(device).unsqueeze(0)
with torch.no_grad():
output = style_transfer(vgg, decoder, mcc_module, content, style,
alpha)
output = output.squeeze(0)
return output.cpu()
content = content.to(device)
with torch.no_grad():
output = style_transfer(vgg, decoder, content, style, args.alpha,
interpolation_weights)
output = output.cpu()
output_name = '{:s}/{:s}_interpolation{:s}'.format(
args.output,
splitext(basename(content_path))[0], args.save_ext)
save_image(output, output_name)
else: # process one content and one style
for style_path in style_paths:
content = content_tf(Image.open(content_path).convert('RGB'))
style = style_tf(Image.open(style_path).convert('RGB'))
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, content, style,
args.alpha)
output = F.upsample(output, size=(256, 256), mode='nearest')
output = output.cpu()
output_name = '{:s}/{:s}_stylized_{:s}{:s}'.format(
args.output,
splitext(basename(content_path))[0],
splitext(basename(style_path))[0], args.save_ext)
save_image(output, output_name)
def make_output():
content_tf = test_transform(args.content_size, args.crop)
style_tf = test_transform(args.style_size, args.crop)
root = Path(
"/HDD2/amirul/datasets/pascal_2012_semantic_segmentation/VOCdevkit/VOC2012"
)
img_root = root / "JPEGImages"
dst_root = root / "style_trans"
if not dst_root.is_dir():
dst_root.mkdir()
style_dict = {}
for k, v in STYLES.items():
style = style_tf(Image.open(v))
if args.preserve_color:
style = coral(style, content)
style = style.to(device).unsqueeze(0)
style_dict[k] = style
train_file = "/mnt/zeta_share_1/public_share/neuralstyle_weakly/irn/voc12/train_aug.txt"
with open(train_file, 'r') as reader:
lines = reader.readlines()
lines = map(lambda x: x.rstrip("\n"), lines)
content_paths = [img_root / f for f in lines]
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)
for f in content_paths:
content = content_tf(Image.open(f.with_suffix(".jpg")))
ori = content.clone().permute(1, 2, 0)
ori = toIMG(ori)
out_name = (dst_root / (f.name + "_ori")).with_suffix(".jpg")
#sio.imsave(out_name, ori)
content = content.to(device).unsqueeze(0)
for style_name, style in style_dict.items():
with torch.no_grad():
output = style_transfer(vgg, decoder, content, style,
args.alpha)
output = output.squeeze().permute(1, 2, 0).cpu()
output = toIMG(output)
out_name = (dst_root /
(f.name + "_" + style_name)).with_suffix(".jpg")
#sio.imsave(out_name, output)
plt.imshow(output)
plt.show()
"""
thumbnail = image.resize(
(int(aspect_ratio * args.thumb_size), args.thumb_size))
thumbnail = content_tf(thumbnail).to(device)
patches = preprocess(image,
patch_size=PATCH_SIZE,
padding=PADDING,
transform=content_tf,
cuda=False)
style = style_tf(style).unsqueeze(0).to(device)
print("patch:", patches.shape)
print("thumbnail:", thumbnail.shape)
print("style:", style.shape)
if args.preserve_color:
style = coral(style, thumbnail)
with torch.no_grad():
style_f = vgg(style)
style_transfer_thumbnail(
thumbnail,
style_f,
save=False if args.test_speed else True,
save_path=os.path.join(args.outf,
"thumb-%d.jpg" % args.thumb_size))
style_transfer_high_resolution(
patches,
style_f,
padding=PADDING,
collection=False,
save_path=os.path.join(
