text_plain_output_1.png
from PIL import Image

import matplotlib.pyplot as plt

import torch

import torchvision.transforms as transforms

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

imsize = (512, 220) if torch.cuda.is_available() else (128, 220)
loader = transforms.Compose([transforms.Resize(imsize), transforms.ToTensor()])
def image_loader(image_name):
    image = Image.open(image_name)
    image = loader(image).unsqueeze(0)
    return image.to(device, torch.float)
style_img = image_loader('./style.jpg')
content_img = image_loader('./content.jpg')
assert style_img.size() == content_img.size(), 'we need to import style and content images of the same size'

unloader = transforms.ToPILImage()
plt.ion()
def imshow(tensor, title=None):
    image = tensor.cpu().clone()
    image = image.squeeze(0)
    image = unloader(image)
    plt.pause(0.001)

input_img = content_img.clone()
plt.figure()
imshow(input_img, title='Input Image')