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] % 5 == 0:
print("run {}:".format(run))
print('Style Loss : {:4f} Content Loss: {:4f}'.format(
style_score.item(), content_score.item()))
print()
if prev:
img_handler.imshow(input_img,
unloader,
title='Run {}'.format(run))
if runs:
run_save = os.path.join(run_path, f'run{run}.jpg')
img_handler.imsave(run_save, unloader, input_img)
return style_score + content_score
output_path = os.path.join(output_path, out_img)
print(f'Output path: {output_path}')
# add the original input image to the figure:
if plt_prvs:
plt.figure()
img_handler.imshow(input_img, unloader, title='Input Image')
# --- RUNNING THE ALGORITHM ---
style_model = StyleModel(device,
cnn,
cnn_normalization_mean,
cnn_normalization_std,
style_img,
content_img,
style_layers=style_layers_default,
content_layers=content_layers_default)
output = style_model.run_style_transfer(input_img=input_img,
unloader=unloader,
prev=plt_prvs,
runs=run_saves)
# --- SHOW AND SAVE OUTPUT ---
if plt_prvs:
img_handler.imshow(output, unloader, title='Output Image')
plt.ioff()
plt.show()
img_handler.imsave(output_path, unloader, output)
input("Press enter to continue...")
