以下是将生成器改造成能够匹配edge-connect中的InpaintingModel的预训练模型键值的结构的代码：
import torch
import torch.nn as nn
import cv2 as cv
from torchvision import transforms
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.encoder = nn.Sequential(
nn.Conv2d(3, 64, 3, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),
nn.Conv2d(64, 128, 3, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
nn.Conv2d(128, 256, 3, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
nn.Conv2d(256, 512, 3, stride=2, padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),
nn.Conv2d(512, 4000, 1),
nn.BatchNorm2d(4000),
nn.LeakyReLU(0.2)
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(4000, 512, 3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),
nn.ConvTranspose2d(512, 256, 3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),
nn.ConvTranspose2d(64, 3, 3, stride=1, padding=1),
nn.Tanh()
)
def forward(self, x, mask):
x = x * (1 - mask)
x = self.encoder(x)
x = self.decoder(x)
x = x * (1 - mask) + x * mask
return x
def load_edgeconnect_weights(model, weight_path):
state_dict = torch.load(weight_path)
new_state_dict = {}
for key, value in state_dict.items():
if 'netG.' in key:
new_key = key.replace('netG.', '')
new_state_dict[new_key] = value
model.load_state_dict(new_state_dict)
# 读取图像和遮罩
img = cv.imread('example.jpg')[:, :, ::-1] / 255.0
mask = cv.inRange(img, (0, 0, 0), (1, 1, 1))
# 转换为张量
image_tensor = transforms.ToTensor()(img)
mask_tensor = transforms.ToTensor()(mask)
# 扩展维度
image_tensor = image_tensor.unsqueeze(0)
mask_tensor = mask_tensor.unsqueeze(0)
# 加载预训练模型权重
generator = Generator()
load_edgeconnect_weights(generator, 'gen.pth')
# 将张量和模型移动到GPU上
image_tensor = image_tensor.cuda()
mask_tensor = mask_tensor.cuda()
generator = generator.cuda()
# 使用预训练模型生成修复后的图像
with torch.no_grad():
output_tensor = generator(image_tensor, mask_tensor)
# 将张量转换为图像
output = output_tensor.detach().cpu().squeeze().numpy().transpose(1, 2, 0)
output = cv.cvtColor(output, cv.COLOR_RGB2BGR)
cv.imwrite('output.jpg', output * 255)
1388
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
