1. 合成图像与真实图像风格迁移基础

合成图像与真实图像风格迁移基础 学习卡片

任务说明

本案例实现一个最小可用的CycleGAN模型，完成合成图像到真实图像的风格迁移。输入为合成图像（如渲染的机械零件图），输出为风格接近真实拍摄图像的逼真图像。模型包含两个生成器（G：合成→真实，F：真实→合成）和两个判别器（Dreal，Dsyn），使用循环一致性损失约束内容保持。最终通过FID（Fréchet Inception Distance）和人工目视检查评估迁移效果。

模型架构与数据准备

生成器采用ResNet-9块结构的卷积网络，判别器采用70×70 PatchGAN。训练数据使用小型合成图像集（如ShapeNet渲染图）与真实图像集（如真实零件照片），各100张，尺寸统一为128×128。以下代码定义完整可训练的CycleGAN，包含必要的初始化、损失计算和训练循环。


import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder
import os
from PIL import Image

# ---------- 生成器：ResNet-9块 ----------
class ResidualBlock(nn.Module):
    def __init__(self, in_channels):
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.InstanceNorm2d(in_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels, in_channels, 3, padding=1),
            nn.InstanceNorm2d(in_channels)
        )
    def forward(self, x):
        return x + self.block(x)

class Generator(nn.Module):
    def __init__(self, in_channels=3, out_channels=3):
        super().__init__()
        # 下采样
        self.down1 = nn.Sequential(nn.Conv2d(in_channels, 64, 7, padding=3), nn.InstanceNorm2d(64), nn.ReLU(inplace=True))
        self.down2 = nn.Sequential(nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.InstanceNorm2d(128), nn.ReLU(inplace=True))
        self.down3 = nn.Sequential(nn.Conv2d(128, 256, 3, stride=2, padding=1), nn.InstanceNorm2d(256), nn.ReLU(inplace=True))
        # ResNet块
        self.res_blocks = nn.Sequential(*[ResidualBlock(256) for _ in range(9)])
        # 上采样
        self.up1 = nn.Sequential(nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1), nn.InstanceNorm2d(128), nn.ReLU(inplace=True))
        self.up2 = nn.Sequential(nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1), nn.InstanceNorm2d(64), nn.ReLU(inplace=True))
        self.out = nn.Sequential(nn.Conv2d(64, out_channels, 7, padding=3), nn.Tanh())
    def forward(self, x):
        x = self.down1(x); x = self.down2(x); x = self.down3(x)
        x = self.res_blocks(x)
        x = self.up1(x); x = self.up2(x)
        return self.out(x)

# ---------- 判别器：70×70 PatchGAN ----------
class Discriminator(nn.Module):
    def __init__(self, in_channels=3):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_channels, 64, 4, stride=2, padding=1), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, 4, stride=2, padding=1), nn.InstanceNorm2d(128), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, 4, stride=2, padding=1), nn.InstanceNorm2d(256), nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 1, 4, padding=1)  # 输出Patch图
        )
    def forward(self, x):
        return self.model(x)

# ---------- 数据集加载 ----------
class PairedDataset(Dataset):
    """从两个独立目录加载合成与真实图像，按文件名顺序配对"""
    def __init__(self, syn_dir, real_dir, size=128):
        self.syn_paths = sorted([os.path.join(syn_dir, f) for f in os.listdir(syn_dir) if f.endswith(('.png','.jpg'))])
        self.real_paths = sorted([os.path.join(real_dir, f) for f in os.listdir(real_dir) if f.endswith(('.png','.jpg'))])
        assert len(self.syn_paths) == len(self.real_paths), "合成与真实图像数量必须一致"
        self.transform = transforms.Compose([
            transforms.Resize((size, size)),
            transforms.ToTensor(),
            transforms.Normalize([0.5]*3, [0.5]*3)
        ])
    def __len__(self):
        return len(self.syn_paths)
    def __getitem__(self, idx):
        syn = self.transform(Image.open(self.syn_paths[idx]).convert('RGB'))
        real = self.transform(Image.open(self.real_paths[idx]).convert('RGB'))
        return syn, real

# ---------- 训练配置 ----------
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size = 4
epochs = 50
lambda_cycle = 10.0

# 初始化模型
G = Generator().to(device)   # 合成→真实
F = Generator().to(device)   # 真实→合成
D_real = Discriminator().to(device)  # 判别真实图像
D_syn = Discriminator().to(device)   # 判别合成图像

# 损失与优化器
criterion_GAN = nn.MSELoss()
criterion_cycle = nn.L1Loss()
optimizer_G = optim.Adam(list(G.parameters()) + list(F.parameters()), lr=0.0002, betas=(0.5, 0.999))
optimizer_D_real = optim.Adam(D_real.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_D_syn = optim.Adam(D_syn.parameters(), lr=0.0002, betas=(0.5, 0.999))

# 加载数据（需准备syn_images/和real_images/两个目录，各含相同数量的图片）
dataset = PairedDataset('syn_images', 'real_images')
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

# ---------- 训练循环 ----------
for epoch in range(epochs):
    for i, (syn, real) in enumerate(dataloader):
        syn, real = syn.to(device), real.to(device)
        batch = syn.size(0)
        valid = torch.ones((batch, 1, 16, 16), device=device)  # PatchGAN标签
        fake = torch.zeros((batch, 1, 16, 16), device=device)

        # 训练生成器
        optimizer_G.zero_grad()
        fake_real = G(syn)          # 合成→真实
        rec_syn = F(fake_real)      # 循环重建合成
        fake_syn = F(real)          # 真实→合成
        rec_real = G(fake_syn)      # 循环重建真实

        loss_gan_G = criterion_GAN(D_real(fake_real), valid) + criterion_GAN(D_syn(fake_syn), valid)
        loss_cycle = criterion_cycle(rec_syn, syn) + criterion_cycle(rec_real, real)
        loss_G = loss_gan_G + lambda_cycle * loss_cycle
        loss_G.backward()
        optimizer_G.step()

        # 训练判别器D_real
        optimizer_D_real.zero_grad()
        loss_D_real = criterion_GAN(D_real(real), valid) + criterion_GAN(D_real(fake_real.detach()), fake)
        loss_D_real.backward()
        optimizer_D_real.step()

        # 训练判别器D_syn
        optimizer_D_syn.zero_grad()
        loss_D_syn = criterion_GAN(D_syn(syn), valid) + criterion_GAN(D_syn(fake_syn.detach()), fake)
        loss_D_syn.backward()
        optimizer_D_syn.step()

    if (epoch+1) % 10 == 0:
        print(f'Epoch [{epoch+1}/{epochs}] Loss_G: {loss_G.item():.4f}, Loss_D_real: {loss_D_real.item():.4f}, Loss_D_syn: {loss_D_syn.item():.4f}')

# 保存模型
torch.save(G.state_dict(), 'generator_syn2real.pth')
print('训练完成，模型已保存为 generator_syn2real.pth')

推理与风格迁移

训练完成后，加载生成器G对任意合成图像进行推理，输出即为风格迁移后的逼真图像。以下代码演示单张图像的迁移过程：


def inference(model_path, input_image_path, output_image_path):
    model = Generator().to(device)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    transform = transforms.Compose([
        transforms.Resize((128, 128)),
        transforms.ToTensor(),
        transforms.Normalize([0.5]*3, [0.5]*3)
    ])
    img = transform(Image.open(input_image_path).convert('RGB')).unsqueeze(0).to(device)
    with torch.no_grad():
        fake = model(img)
    # 反归一化并保存
    fake = (fake.squeeze(0).cpu().permute(1,2,0).numpy() * 0.5 + 0.5) * 255
    fake = fake.astype('uint8')
    Image.fromarray(fake).save(output_image_path)
    print(f'迁移结果已保存至 {output_image_path}')

# 使用示例（取消注释并替换路径后执行）
# inference('generator_syn2real.pth', 'test_syn.png', 'test_real_style.png')

验证方法与判定标准

运行方式：准备两个目录syn_images/和real_images/，各放入100张128×128以上的图像（文件名按相同顺序排列），执行训练脚本。训练完成后运行推理函数处理测试图像。

判定标准：

  1. 训练收敛：观察打印的Loss值，生成器损失应逐渐下降并趋于稳定，判别器损失在0.5附近波动（表示对抗平衡）；
  2. 视觉质量：目视检查输出图像，应保留合成图像的内容结构（形状、位置），但纹理、光照、颜色分布接近真实图像风格；
  3. 定量指标：使用FID评估（需额外安装pytorch-fid库），计算100张合成图像迁移前后与真实图像集的FID。预期迁移后FID显著低于迁移前（例如从80+降至50以下），说明分布更接近真实图像。
CycleGAN 合成→真实风格迁移架构 合成图像 X_syn 生成器 G 合成→真实 伪真实图像 Y_fake = G(X_syn) 判别器 D_real 真/伪 循环一致性损失 L_cycle 生成器 F 真实→合成 真实图像 Y_real 伪合成图像 X_fake = F(Y_real) 判别器 D_syn 真/伪 循环一致性损失 L_cycle 生成器 判别器 图像域 循环一致性路径

关键实现说明

本实现采用最小完整CycleGAN结构：生成器使用9个残差块保持图像内容，InstanceNorm保证风格迁移的稳定性；判别器使用PatchGAN输出16×16的判别图，增强局部纹理判别能力。训练时合成与真实图像按文件名顺序配对，实际应用中无需严格配对，CycleGAN天然支持非配对数据。若需处理非配对数据，只需将PairedDataset改为分别从两个目录随机采样的方式即可，不影响模型主体。

← 上一章 📖 返回目录 下一章 →