def run_style_transfer(content_image,
style_image,
input_image,
num_steps=300,
style_weight=1000000,
content_weight=1,
path=""):
vgg_model, style_losses, content_losses = get_style_model_and_losses(
style_image,
content_image,
style_layers=style_layers,
content_layers=content_layers)
optimizer = get_input_optimizer(input_image)
i = [0]
input_image.data.clamp_(0, 1)
if path != "":
path += "\\{}.jpg"
save(input_image, path.format(i[0]))
while i[0] <= num_steps:
def closure():
optimizer.zero_grad()
vgg_model(input_image)
style_score = 0
content_score = 0
for s in style_losses:
style_score += s.loss
for c in content_losses:
content_score += c.loss
style_score *= style_weight
content_score *= content_weight
loss = style_score + content_score
loss.backward()
# puede darse que la imagen haya acabado con píxeles fuera de [0,1], esto lo arregla
input_image.data.clamp_(0, 1)
i[0] += 1
if i[0] % 50 == 0:
print("iteraciones: {}:".format(i[0]))
print(
'Diferencia de estilo : {:4f} Diferencia de contenido: {:4f}'
.format(style_score.item(), content_score.item()))
if path != "":
save(input_image, path.format(i[0]))
print()
return style_score + content_score
optimizer.step(closure)
return input_image
def run_style_transfer(cnn,
normalization_mean,
normalization_std,
content_img,
style_img,
input_img,
num_steps=300,
style_weight=1000000,
content_weight=1):
"""Run the style transfer."""
print('Building the style transfer model..')
model, style_losses, content_losses = get_style_model_and_losses(
cnn, normalization_mean, normalization_std, style_img, content_img)
optimizer = get_input_optimizer(input_img)
print('Optimizing..')
run = [0]
while run[0] <= num_steps:
def closure():
# correct the values of updated input image
input_img.data.clamp_(0, 1)
optimizer.zero_grad()
model(input_img)
style_score = 0
content_score = 0
for sl in style_losses:
style_score += sl.loss
for cl in content_losses:
content_score += cl.loss
style_score *= style_weight
content_score *= content_weight
loss = style_score + content_score
loss.backward()
run[0] += 1
if run[0] % 50 == 0:
print("run {}:".format(run))
print('Style Loss : {:4f} Content Loss: {:4f}'.format(
style_score.item(), content_score.item()))
print()
return style_score + content_score
optimizer.step(closure)
# a last correction...
input_img.data.clamp_(0, 1)
return input_img
# if you want to use white noise instead uncomment the below line:
# input_img = torch.randn(content_img.data.size(), device=device)
style_mask, content_mask = load_masks(style_img.size(2), style_img.size(3))
cnn = models.vgg19(pretrained=True).features.to(device).eval()
cnn_normalization_mean = torch.tensor([0.485, 0.456, 0.406]).to(device)
cnn_normalization_std = torch.tensor([0.229, 0.224, 0.225]).to(device)
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
print('Building the style transfer model..')
model, style_losses, content_losses, sim_losses = get_style_model_and_losses(
cnn, cnn_normalization_mean, cnn_normalization_std, style_img,
content_img, style_mask, content_mask, device)
optimizer = optim.LBFGS([input_img.requires_grad_()], lr=args.lr)
print('Optimizing..')
run = [0]
while run[0] <= args.iters:
def closure():
optimizer.zero_grad()
model(input_img)
style_score = 0
content_score = 0
sim_score = 0
def run_style_transfer(args):
""" Run the style transfer """
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Building the style transfer model..')
model = models.vgg19(pretrained=True)
style_img = image_loader(args.style_img, device)
content_img = image_loader(args.content_img, device)
input_img = get_input_img(content_img, device, use_noise=args.use_noise)
anim_lst = []
model, style_losses, content_losses = get_style_model_and_losses(
model, style_img, content_img, device)
optimizer = get_input_optimizer(input_img)
print('Optimizing..')
run = 0
while run < args.num_steps:
def closure():
"""
Some optimization algorithms such as Conjugate Gradient and LBFGS need to reevaluate
the function multiple times, so you have to pass in a closure that allows them to
recompute your model. The closure should clear the gradients, compute the loss, and
return it. (Refer to https://pytorch.org/docs/stable/optim.html)
Update the parameters by optimizer.step(closure)
"""
input_img.data.clamp_(
0, 1) # correct the values of updated input image
nonlocal run # enable rebinding of a nonlocal name
# https://stackoverflow.com/questions/2609518/unboundlocalerror-with-nested-function-scopes
if args.save_anim and run % 10 == 0:
anim_lst.append(input_img.cpu().clone())
optimizer.zero_grad()
model(input_img)
style_score = 0
content_score = 0
for sl in style_losses:
style_score += sl.loss
for cl in content_losses:
content_score += cl.loss
style_score *= args.style_weight
content_score *= args.content_weight
total_loss = style_score + content_score
total_loss.backward()
run += 1
if run % 25 == 0:
print("run {}".format(run))
print("Style Loss: {:4f} Content Loss: {:4f}".format(
style_score.item(), content_score.item()))
print()
return total_loss
optimizer.step(closure)
input_img.data.clamp_(0, 1)
if args.save_anim:
anim_lst.append(input_img.cpu().clone())
animate(anim_lst, 'images\output.gif')
show_image(input_img,
title='Output Image',
save=True,
save_path=args.output_path)