简易AI绘图python代码

bingwangzi · 发表于 2023-11-5 11:15:10

import torch
import torchvision.transforms as transforms
from PIL import Image
from stylegan2 import StyleGAN2

# 加载模型
model = StyleGAN2()
model.load_state_dict(torch.load('stylegan2.pth'))
model.eval()

# 定义图像变换
transform = transforms.Compose([
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
])

# 加载图像
image = Image.open('input.jpg')
image = transform(image)
image = image.unsqueeze(0) # 添加 batch 维度

# 进行风格迁移
with torch.no_grad():
output = model(image)
output = output.squeeze(0).permute(1, 2, 0) # 去除 batch 维度，并转回正常通道顺序

# 保存输出图像
output_image = Image.fromarray((output * 0.5 + 0.5).clamp(0, 1).byte().cpu().numpy())
output_image.save('output.jpg')

上述代码中的 stylegan2.pth 文件需要预先下载并放置在代码所在的目录下，该文件包含了 StyleGAN2 的模型权重。另外，输入图像 input.jpg 也需要预先准备好并放置在代码所在的目录下。输出图像将保存为 output.jpg。

cjcjwokan · 发表于 2023-11-24 11:24:52

看不懂，非常不懂

zzj · 发表于 2023-12-5 08:45:06

看不懂看不懂

yang67 · 发表于 2024-1-10 00:04:43

谢谢楼主分享.

zhanggangyu · 发表于 2024-2-29 21:20:03

感谢分享

		自动登录	找回密码
密码			立即注册