给出edgeconnec中InpaintingModel的生成器原版代码和使用InpaintingModel_gen.pth预训练模型进行图像修复的代码

InpaintingModel的生成器原版代码如下：

import torch.nn as nn
import torch.nn.functional as F

class InpaintGenerator(nn.Module):
    def __init__(self, image_channels):
        super().__init__()
        
        self.conv1 = nn.Conv2d(image_channels, 64, kernel_size=6, stride=2, padding=2)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1)
        self.conv4 = nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1)
        self.conv5 = nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.conv6 = nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.conv7 = nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.conv8 = nn.Conv2d(512, 512, kernel_size=4, stride=2, padding=1)
        
        self.convT1 = nn.ConvTranspose2d(512, 512, kernel_size=4, stride=2, padding=1)
        self.convT2 = nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1)
        self.convT3 = nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1)
        self.convT4 = nn.ConvTranspose2d(1024, 512, kernel_size=4, stride=2, padding=1)
        self.convT5 = nn.ConvTranspose2d(1024, 256, kernel_size=4, stride=2, padding=1)
        self.convT6 = nn.ConvTranspose2d(512, 128, kernel_size=4, stride=2, padding=1)
        self.convT7 = nn.ConvTranspose2d(256, 64, kernel_size=4, stride=2, padding=1)
        self.convT8 = nn.ConvTranspose2d(128, image_channels, kernel_size=6, stride=2, padding=2)
        
        self.norm1 = nn.InstanceNorm2d(64)
        self.norm2 = nn.InstanceNorm2d(128)
        self.norm3 = nn.InstanceNorm2d(256)
        self.norm4 = nn.InstanceNorm2d(512)
        self.norm5 = nn.InstanceNorm2d(512)
        self.norm6 = nn.InstanceNorm2d(512)
        self.norm7 = nn.InstanceNorm2d(512)
        self.norm8 = nn.InstanceNorm2d(512)
        
        self.normT1 = nn.InstanceNorm2d(512)
        self.normT2 = nn.InstanceNorm2d(512)
        self.normT3 = nn.InstanceNorm2d(512)
        self.normT4 = nn.InstanceNorm2d(512)
        self.normT5 = nn.InstanceNorm2d(256)
        self.normT6 = nn.InstanceNorm2d(128)
        self.normT7 = nn.InstanceNorm2d(64)
        
    def forward(self, x):
        x = F.relu(self.norm1(self.conv1(x)))
        x = F.relu(self.norm2(self.conv2(x)))
        x = F.relu(self.norm3(self.conv3(x)))
        x = F.relu(self.norm4(self.conv4(x)))
        x = F.relu(self.norm5(self.conv5(x)))
        x = F.relu(self.norm6(self.conv6(x)))
        x = F.relu(self.norm7(self.conv7(x)))
        x = F.relu(self.norm8(self.conv8(x)))
        
        x = F.relu(self.normT1(self.convT1(x)))
        x = F.relu(self.normT2(self.convT2(torch.cat([x, self.norm7(self.conv7(x))], dim=1))))
        x = F.relu(self.normT3(self.convT3(torch.cat([x, self.norm6(self.conv6(x))], dim=1))))
        x = F.relu(self.normT4(self.convT4(torch.cat([x, self.norm5(self.conv5(x))], dim=1))))
        x = F.relu(self.normT5(self.convT5(torch.cat([x, self.norm4(self.conv4(x))], dim=1))))
        x = F.relu(self.normT6(self.convT6(torch.cat([x, self.norm3(self.conv3(x))], dim=1))))
        x = F.relu(self.normT7(self.convT7(torch.cat([x, self.norm2(self.conv2(x))], dim=1))))
        x = self.convT8(torch.cat([x, self.norm1(self.conv1(x))], dim=1))
        x = F.tanh(x)
        return x

使用InpaintingModel_gen.pth预训练模型进行图像修复的代码如下：

import torch
from edgeconnect.models import InpaintGenerator

# 加载预训练模型
model = InpaintGenerator(image_channels=3)
model.load_state_dict(torch.load('InpaintingModel_gen.pth', map_location=torch.device('cpu')))
model.eval()

# 读取待修复的图像
img = Image.open('image.png')
img = transforms.ToTensor()(img).unsqueeze(0)

# 进行图像修复
with torch.no_grad():
    output = model(img)

# 保存修复后的图像
output = output.squeeze().permute(1,2,0).cpu().numpy()
output = (output + 1) / 2.0 * 255.0
output = output.astype(np.uint8)
Image.fromarray(output).save('output.png')