用上下文编码加GAN修复128*128带有随机掩膜区域的图像，提供掩膜区域像素设置为0 的代码

代码如下：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.utils.data import DataLoader
import numpy as np
from models import Generator, Discriminator

# 设定超参数
batch_size = 64  # mini-batch大小
image_size = 128  # 图像大小
nz = 100  # 生成器输入噪声维度
ngf = 64  # 生成器特征图大小
ndf = 64  # 判别器特征映射大小
num_epochs = 200  # 训练轮数
lr = 0.0002  # 学习率
beta1 = 0.5   # Adam优化器的参数

# 加载数据集
transform = transforms.Compose([
    transforms.Resize(image_size),
    transforms.CenterCrop(image_size),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_dataset = dset.ImageFolder(root='./data', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

# 创建生成器和判别器
generator = Generator(nz, ngf)
discriminator = Discriminator(ndf)

# 定义损失函数和优化器
criterion = nn.BCELoss()
optimizer_g = optim.Adam(generator.parameters(), lr=lr, betas=(beta1, 0.999))
optimizer_d = optim.Adam(discriminator.parameters(), lr=lr, betas=(beta1, 0.999))

# 定义一些辅助函数
def mask_random(img):
    # 生成随机掩膜
    mask = np.ones((image_size, image_size, 3))
    c, h, w = img.size()
    x = np.random.randint(image_size-h)
    y = np.random.randint(image_size-w)
    mask[x:x+h, y:y+w, :] = 0
    img_masked = img * torch.from_numpy(mask).float()
    return img_masked, mask

# 进行训练
for epoch in range(num_epochs):
    for i, (imgs, _) in enumerate(train_loader):
        batch_size = imgs.size(0)
        
        # 真实图像的标签为1，假图像的标签为0
        real_labels = torch.ones(batch_size)
        fake_labels = torch.zeros(batch_size)
        
        # 提取掩膜部分图像
        imgs_masked = []
        masks = []
        for img in imgs:
            img_masked, mask = mask_random(img)
            imgs_masked.append(img_masked)
            masks.append(torch.from_numpy(mask).permute(2, 0, 1))
        imgs_masked = torch.stack(imgs_masked)
        masks = torch.stack(masks)
        
        # 训练判别器
        discriminator.zero_grad()
        real_outputs = discriminator(imgs).squeeze()
        fake_imgs = generator(torch.randn(batch_size, nz, 1, 1))
        fake_outputs = discriminator(fake_imgs.detach()).squeeze()
        d_loss = criterion(real_outputs, real_labels) + criterion(fake_outputs, fake_labels)
        d_loss.backward()
        optimizer_d.step()
        
        # 训练生成器
        generator.zero_grad()
        fake_imgs = generator(torch.randn(batch_size, nz, 1, 1))
        fake_outputs = discriminator(fake_imgs).squeeze()
        g_loss = criterion(fake_outputs, real_labels)
        
        # 预测掩膜部分图像
        img_predicted = imgs_masked.clone()
        img_predicted[:, :, x:x+h, y:y+w] = fake_imgs[:, :, x:x+h, y:y+w]
        
        # 计算修复图像和真实图像的损失
        img_repaired = img_predicted * masks + imgs * (1 - masks)
        l1_loss = nn.L1Loss()(img_repaired, imgs)
        g_total_loss = g_loss + 10 * l1_loss
        
        g_total_loss.backward()
        optimizer_g.step()
        
        # 输出日志
        if i % 50 == 0:
            print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f L1_Loss: %.4f'
                  % (epoch, num_epochs, i, len(train_loader),
                     d_loss.item(), g_loss.item(), l1_loss.item()))
            
        # 每50个batch可视化一组结果
        if i % 50 == 0:
            vutils.save_image(imgs, './results/real_samples.png')
            fake_imgs = generator(torch.randn(batch_size, nz, 1, 1)).detach()
            fake_imgs[:, :, x:x+h, y:y+w] = img_masked[:, :, x:x+h, y:y+w]
            vutils.save_image(fake_imgs, './results/fake_samples_epoch_%03d_batch_%03d.png' % (epoch, i))

提示：在代码中，我们使用了一个由生成器生成的修复图像，并计算该修复图像与真实图像之间的L1损失，以促进生成器生成质量更高的修复结果。通过使用该技术，我们可以使生成器不仅仅是复制掩膜区域的周围像素，而是学习更全面的图像特征以有效地修复缺失的像素。