|
|
from torchvision import utils as vutils
|
|
|
|
|
|
from models import *
|
|
|
from utils import *
|
|
|
|
|
|
CORE = "cuda:0" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
device = torch.device(CORE)
|
|
|
|
|
|
transform_net = TransformNet(32).to(device)
|
|
|
transform_net.load_state_dict(
|
|
|
torch.load('./model/Picasso.pth', map_location=CORE))
|
|
|
|
|
|
style_img = read_image('./style/style-Picasso.jpg').to(device)
|
|
|
|
|
|
content_img = read_image('./img/content.jpg').to(device)
|
|
|
output_img = transform_net(content_img)
|
|
|
|
|
|
plt.figure(figsize=(18, 6))
|
|
|
|
|
|
plt.subplot(1, 3, 1)
|
|
|
imshow(style_img, title='Style Image')
|
|
|
|
|
|
plt.subplot(1, 3, 2)
|
|
|
imshow(content_img, title='Content Image')
|
|
|
|
|
|
plt.subplot(1, 3, 3)
|
|
|
imshow(output_img.detach(), title='Output Image')
|
|
|
|
|
|
plt.show()
|
|
|
|
|
|
# def save_image_tensor2cv2(input_tensor: torch.Tensor, filename):
|
|
|
# """
|
|
|
# 将tensor保存为cv2格式
|
|
|
# :param input_tensor: 要保存的tensor
|
|
|
# :param filename: 保存的文件名
|
|
|
# """
|
|
|
# assert (len(input_tensor.shape) == 4 and input_tensor.shape[0] == 1)
|
|
|
# # 复制一份
|
|
|
# input_tensor = input_tensor.clone().detach()
|
|
|
# # 到cpu
|
|
|
# input_tensor = input_tensor.to(torch.device('cpu'))
|
|
|
# # 反归一化
|
|
|
# # input_tensor = unnormalize(input_tensor)
|
|
|
# # 去掉批次维度
|
|
|
# input_tensor = input_tensor.squeeze()
|
|
|
# # 从[0,1]转化为[0,255]，再从CHW转为HWC，最后转为cv2
|
|
|
# input_tensor = input_tensor.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).type(torch.uint8).numpy()
|
|
|
# # RGB转BRG
|
|
|
# input_tensor = cv2.cvtColor(input_tensor, cv2.COLOR_RGB2BGR)
|
|
|
# cv2.imwrite(filename, input_tensor)
|
|
|
|
|
|
vutils.save_image(output_img.detach(), './result/Picasso.jpg')
|
