基于CycleGAN的图像风格转换基于CycleGAN的图像风格转换

CycleGAN（循环生成对抗网络）是一种用于图像到图像转换的深度学习技术，它能够在没有成对训练样本的情况下，将一种风格的图像转换成另一种风格。CycleGAN 通常用于图像风格迁移、季节转换、艺术风格模仿等任务。它是由朱俊彦等人提出的，并在论文《Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks》中进行了详细描述。

CycleGAN 的工作原理如下：

生成器（Generator）：CycleGAN 有两个生成器，分别用于将图像从源域转换到目标域，以及将图像从目标域转换回源域。生成器的目标是生成足够真实的图像，以欺骗判别器。

判别器（Discriminator）：CycleGAN 也有两个判别器，分别用于判断图像是否属于源域或目标域。判别器的目标是能够准确地区分真实图像和生成器生成的假图像。
循环一致性损失（Cycle Consistency Loss）：这是 CycleGAN 的核心概念之一。循环一致性损失确保图像经过一次完整的转换循环（即从源域转换到目标域，再从目标域转换回源域）后，能够回到原始图像。这种损失函数有助于在没有成对数据的情况下保持图像内容的一致性。

身份损失（Identity Loss）：这个损失函数确保当源域图像输入到目标域的生成器时，输出图像应该与输入图像尽可能相似。这样可以避免生成器在转换过程中改变不应该改变的内容。

对抗损失（Adversarial Loss）：这是生成对抗网络（GAN）中常用的损失函数，用于让生成器生成的图像更接近于真实图像。
通过最小化这些损失函数，CycleGAN 能够学习到源域和目标域之间的映射关系，从而实现图像风格的转换。

CycleGAN 的应用非常广泛，例如：

将普通照片转换为印象派风格的画作。
将夏天的风景照片转换为冬天的风景。
将马的图像转换为斑马的图像。

由于CycleGAN 不需要成对的训练数据，它为图像风格转换提供了一种灵活且强大的工具。

from random import randint
import numpy as np 
import torch
torch.set_default_tensor_type(torch.FloatTensor)
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms
import os
import matplotlib.pyplot as plt
import torch.nn.functional as F
from torch.autograd import Variable
from torchvision.utils import save_image
import shutil
import cv2
import random
from PIL import Image
import itertools

导入代码中需要的各种库和模块，包括PyTorch及其相关库，以及Python的标准库和其他图像处理库。

def to_img(x):
    # 将归一化后的图像转换为0-1范围
    out = 0.5 * (x + 1)
    out = out.clamp(0, 1)  
    # 改变图像的形状以适应显示和保存
    out = out.view(-1, 3, 256, 256)  
    return out

将归一化后的图像数据转换回0到1的范围，并调整其形状以适应图像显示和保存。

数据加载

data_path = os.path.abspath('../data')
image_size = 256
batch_size = 1

transform = transforms.Compose([transforms.Resize(int(image_size * 1.12), 
                                                  Image.BICUBIC), 
            transforms.RandomCrop(image_size), 
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))])

def _get_train_data(batch_size=1):
    
    train_a_filepath = data_path + '\\trainA\\'
    train_b_filepath = data_path + '\\trainB\\'
    
    train_a_list = os.listdir(train_a_filepath)
    train_b_list = os.listdir(train_b_filepath)
    
    train_a_result = []
    train_b_result = [] 
    
    numlist = random.sample(range(0, len(train_a_list)), batch_size)
    
    for i in numlist:
        a_filename = train_a_list[i]
        a_img = Image.open(train_a_filepath + a_filename).convert('RGB')
        res_a_img = transform(a_img)
        train_a_result.append(torch.unsqueeze(res_a_img, 0))
        
        b_filename = train_b_list[i]
        b_img = Image.open(train_b_filepath + b_filename).convert('RGB')
        res_b_img = transform(b_img)
        train_b_result.append(torch.unsqueeze(res_b_img, 0))
        
    return torch.cat(train_a_result, dim=0), torch.cat(train_b_result, dim=0)

设置了数据加载的相关参数，包括数据路径、图像大小和批量大小，并定义了图像转换操作。_get_train_data 函数用于从数据集中加载一批训练数据。

"""
残差网络block
flectionPad2d: 镜像填充，例如 0 1 填充至 4个数，则 0 1 0 1
InstanceNorm2d: 对单个样本的每一层特征图抽出来一层层求均值、方差然后归一化
"""

class ResidualBlock(nn.Module):
    def __init__(self, in_features):
        super(ResidualBlock, self).__init__()
        self.block_layer = nn.Sequential(
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features),
            nn.ReLU(inplace=True),
            nn.ReflectionPad2d(1),
            nn.Conv2d(in_features, in_features, 3),
            nn.InstanceNorm2d(in_features))
        
    def forward(self, x):
        return x + self.block_layer(x)

生成器

class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
             
        model = [nn.ReflectionPad2d(3), 
                 nn.Conv2d(3, 64, 7), 
                 nn.InstanceNorm2d(64), 
                 nn.ReLU(inplace=True)]

        in_features = 64
        out_features = in_features * 2
        for _ in range(2):
            model += [nn.Conv2d(in_features, out_features, 
                                3, stride=2, padding=1), 
            nn.InstanceNorm2d(out_features), 
            nn.ReLU(inplace=True)]
            in_features = out_features
            out_features = in_features*2

        for _ in range(9):
            model += [ResidualBlock(in_features)]

        out_features = in_features // 2
        for _ in range(2):
            model += [nn.ConvTranspose2d(
                    in_features, out_features, 
                    3, stride=2, padding=1, output_padding=1), 
                nn.InstanceNorm2d(out_features), 
                nn.ReLU(inplace=True)]
            in_features = out_features
            out_features = in_features // 2

        model += [nn.ReflectionPad2d(3), 
                  nn.Conv2d(64, 3, 7), 
                  nn.Tanh()]

        self.gen = nn.Sequential( * model)
        
    def forward(self, x):
        x = self.gen(x)
        return x 

残差块和生成器的类定义，这些类是CycleGAN模型的核心组成部分。

"""
判别器 
1、这里判别器的最后一层是FCN全卷积网络

avg_pool2d：以均值方式池化，以下述代码为例：
input  = torch.randn(10, 3, 128, 128)
m      = Discriminator()
output = m(input)

此时在 avg_pool2d 前 x.size() 为 torch.Size([10, 1, 14, 14])
对 10个[14, 15]的tensor求均值并返回
tensor([[0.1162],
        [0.1298],
        [0.1266],
        [0.1229],
        [0.1085],
        [0.1121],
        [0.1064],
        [0.1044],
        [0.1077],
        [0.1139]], grad_fn=<ViewBackward>)

"""

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.dis = nn.Sequential(
            nn.Conv2d(3, 64, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(64, 128, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(128, 256, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(256, 512, 4, padding=1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            
            nn.Conv2d(512, 1, 4, padding=1))        
        
    def forward(self, x):
        x = self.dis(x)
        return F.avg_pool2d(x, x.size()[2:]).view(x.size()[0], -1)

判别器的类定义，判别器用于判断输入的图像是真实的还是由生成器生成的。

"""
训练神经网络
存在生成网络 G_A2B 从A类图片生成B类图片
存在生成网络 G_B2A 从B类图片生成A类图片 
存在判别器   D_A   对A类图片真伪进行判断
存在判别器   D_B   对B类图片真伪进行判断 

存在损失函数：criterion_GAN MSELoss：均方损失函数，即 两个向量各分量差的平方
存在损失函数：criterion_cycle L1Loss：平均均对误差，即 两个向量各分量差的绝对值求和再除分量数
存在损失函数: criterion_identity L1Loss:

第一步：训练生成器
1、利用 G_A2B 将 real_B 生成 same_B,获取 real_B 和 same_B 之间的损失
2、利用 G_B2A 将 real_A 生成 same_A,获取 real_A 和 same_A 之间的损失
3、利用 G_A2B 将 real_A 生成 fake_B, 通过 netD_B 判断 fake_B，获取其与判真之间的损失
4、利用 G_B2A 将 real_B 生成 faka_A，通过 netD_A 判断 fake_A，获取其与判真之间的损失
5、利用 G_B2A 将 fake_B 生成 recovered_A，获取其与 real_A 之间的损失
6、利用 G_A2B 将 fake_A 生成 recovered_B, 获取其与 real_B 之间的损失

此时对所有损失求和后训练生成器，此时：

对于 G_A2B：
    1. 如果输入是 B类图片，它生成的图片same_B总体像素层面上接近B
    2. 如果输入是 A类图片，它生成的图片fake_B会具有B类图片的卷积特征
    3. 对于生成的图片fake_B, 它经过 G_B2A 生成的recovered_A 会总体像素层面上接近A
    
    此时我们有了一张 A 经过 G_A2B 和 G_B2A 成为一张 新的A，这张新A像素总体上是A（是一批马）
    但细节纹理上具有 B 的特征（斑马纹理）

同理有 G_B2A

第二步：训练判别器
1、对判别器 netD_A 可以对 real_A 判真 
2、对判别器 netD_A 可以对 fake_A 判伪
3、对判别器 netD_B 可以对 real_B 判真 
4、对判别器 netD_B 可以对 fake_B 判伪
"""

class ReplayBuffer():
    """
    缓存队列，若不足则新增，否则随机替换
    """
    def __init__(self, max_size=50):
        self.max_size = max_size
        self.data = []
        
    def push_and_pop(self, data):
        to_return = []
        for element in data.data:
            element = torch.unsqueeze(element, 0)
            if len(self.data) < self.max_size:
                self.data.append(element)
                to_return.append(element)
            else:
                if random.uniform(0,1) > 0.5:
                    i = random.randint(0, self.max_size-1)
                    to_return.append(self.data[i].clone())
                    self.data[i] = element
                else:
                    to_return.append(element)
        return Variable(torch.cat(to_return))
    

fake_A_buffer = ReplayBuffer()
fake_B_buffer = ReplayBuffer()

netG_A2B = Generator()
netG_B2A = Generator()
netD_A = Discriminator()
netD_B = Discriminator()

criterion_GAN = torch.nn.MSELoss()
criterion_cycle = torch.nn.L1Loss()
criterion_identity = torch.nn.L1Loss()

d_learning_rate = 3e-4  # 3e-4
g_learning_rate = 3e-4
optim_betas = (0.5, 0.999)


g_optimizer = optim.Adam(itertools.chain(netG_A2B.parameters(), 
                                         netG_B2A.parameters()), 
            lr=d_learning_rate)
da_optimizer = optim.Adam(netD_A.parameters(), lr=d_learning_rate)
db_optimizer = optim.Adam(netD_B.parameters(), lr=d_learning_rate)

num_epochs = 1000 

#检查train_b是否有灰度图
"""
print(np.arange(196608).reshape(256,256,3).shape)
train_b_filepath = data_path + "\\trainB\\"
train_b_list = os.listdir(train_b_filepath)
for i in range(len(train_b_list)):
        b_filename = train_b_list[i]
        b_img = Image.open(train_b_filepath + b_filename)
        if np.array(b_img).shape != np.arange(196608).reshape(256,256,3).shape:
            print(b_filename)
            os.remove(train_b_filepath + b_filename)
"""

损失函数和优化器的定义，这些是训练神经网络所必需的。

训练循环

for epoch in range(num_epochs): 

训练循环的开始，代码在这里迭代进行训练，每次迭代称为一个epoch。

    real_a, real_b = _get_train_data(batch_size)
    target_real = torch.full((batch_size,), 1).float()
	target_fake = torch.full((batch_size,), 0).float()

从数据集加载一批真实图像，并创建真实和虚假标签的Tensor。

g_optimizer.zero_grad()

将生成器的梯度清零，为下一次训练做准备。

    # 第一步：训练生成器
    same_B = netG_A2B(real_b).float()
    loss_identity_B = criterion_identity(same_B, real_b) * 5.0   
    same_A = netG_B2A(real_a).float()
    loss_identity_A = criterion_identity(same_A, real_a) * 5.0
    
    fake_B = netG_A2B(real_a).float()
    pred_fake = netD_B(fake_B).float()
    loss_GAN_A2B = criterion_GAN(pred_fake, target_real)
    fake_A = netG_B2A(real_b).float()
    pred_fake = netD_A(fake_A).float()
    loss_GAN_B2A = criterion_GAN(pred_fake, target_real)
    recovered_A = netG_B2A(fake_B).float()
    loss_cycle_ABA = criterion_cycle(recovered_A, real_a) * 10.0
    recovered_B = netG_A2B(fake_A).float()
    loss_cycle_BAB = criterion_cycle(recovered_B, real_b) * 10.0  
    loss_G = (loss_identity_A + loss_identity_B + loss_GAN_A2B + 
              loss_GAN_B2A + loss_cycle_ABA + loss_cycle_BAB)
    loss_G.backward()    
    g_optimizer.step()

生成器训练的代码，包括前向传播、计算损失和反向传播。

    # 第二步：训练判别器
    # 训练判别器A
    da_optimizer.zero_grad()
    pred_real = netD_A(real_a).float()
    loss_D_real = criterion_GAN(pred_real, target_real)
    fake_A = fake_A_buffer.push_and_pop(fake_A)
    pred_fake = netD_A(fake_A.detach()).float()
    loss_D_fake = criterion_GAN(pred_fake, target_fake)
    loss_D_A = (loss_D_real + loss_D_fake) * 0.5
    loss_D_A.backward()
    da_optimizer.step()
    # 训练判别器B
    db_optimizer.zero_grad()
    pred_real = netD_B(real_b)
    loss_D_real = criterion_GAN(pred_real, target_real)
    fake_B = fake_B_buffer.push_and_pop(fake_B)
    pred_fake = netD_B(fake_B.detach())
    loss_D_fake = criterion_GAN(pred_fake, target_fake)
    loss_D_B = (loss_D_real + loss_D_fake) * 0.5
    loss_D_B.backward()
    db_optimizer.step()   

判别器训练的代码，包括前向传播、计算损失和反向传播。

  #损失打印，存储伪造图片
  print('Epoch[{}],loss_G:{:.6f} ,loss_D_A:{:.6f},loss_D_B:{:.6f}'
    .format(epoch, loss_G.data.item(), loss_D_A.data.item(), 
            loss_D_B.data.item()))
  if (epoch + 1) % 20 == 0 or epoch == 0:  
      b_fake = to_img(fake_B.data)
      a_fake = to_img(fake_A.data)
      a_real = to_img(real_a.data)
      b_real = to_img(real_b.data)
      save_image(a_fake, '../tmp/a_fake.png') 
      save_image(b_fake, '../tmp/b_fake.png') 
      save_image(a_real, '../tmp/a_real.png') 
      save_image(b_real, '../tmp/b_real.png') 

再每次迭代结束时，打印出当前的总损失，并在每个epoch存储生成的图像，将生成的伪造图像和真实的训练图像保存到指定的目录中。这有助于可视化训练过程中的进展，并检查生成器是否能够生成逼真的图像。

总体代码结构:
利用PyTorch实现的CycleGAN的模型训练过程。CycleGAN是一种可以学习从一个域转换到另一个域的图像到图像的转换的生成对抗网络，不需要成对的训练数据。
1.导入必要的库：代码首先导入了所需的库，如random, numpy, torch, torchvision等。
2.定义数据加载函数 _get_train_data：这个函数从指定的数据文件夹中随机读取一对图片作为训练数据。
3.定义残差块 ResidualBlock：这是CycleGAN中的一个基本组件，用于构建生成器和判别器。
4.定义生成器 Generator：生成器用于将图片从一个域转换到另一个域。
5.定义判别器 Discriminator：判别器用于判断输入图片是否属于特定域。
6.定义重播缓冲区 ReplayBuffer：用于在训练过程中存储生成的图片，以用于后续的训练。
7.初始化生成器和判别器：创建生成器 netG_A2B, netG_B2A 和判别器 netD_A, netD_B 的实例。
8.定义损失函数：包括对抗损失 criterion_GAN, 循环一致性损失 criterion_cycle 和身份损失 criterion_identity。
9.定义优化器：为生成器和判别器定义优化器。
10.训练循环：在指定数量的 epochs 内进行训练。在每次迭代中，首先更新生成器的参数，然后更新判别器的参数。同时，会定期保存生成的图片和打印损失。
11.损失打印和图片存储：在训练过程中，会定期打印出当前的损失，并保存生成的图片。
总体来说，实现了CycleGAN的训练过程，通过迭代地更新生成器和判别器的参数，使得生成器能够学会从一个域转换到另一个域的图片，而判别器能够准确地判断图片属于哪个域。

运行结果

CycleGAN模型在训练过程中的一个日志记录。显示了当前训练状态的一些关键指标，这些指标对于理解模型的性能和训练进度非常重要。
Epoch[x]: 表示这是训练过程中的第x个epoch。在CycleGAN中，epoch是训练过程的基本单位，通常代表一次完整的数据集遍历。
loss_G:x.xxxxxx: 这是生成器（Generator）的损失。在CycleGAN中，生成器的损失由多个部分组成，包括身份损失（Identity Loss）、对抗损失（GAN Loss）和循环一致性损失（Cycle Consistency Loss）。这里的损失值是一个综合了所有这些损失的数值。

loss_D_A:0.xxxxxx: 这是判别器A（Discriminator A）的损失。判别器A用于判断A类图片的真伪，它的损失也是由多个部分组成的，包括判断真实A类图片的损失和判断生成器生成的虚假A类图片的损失。这里的损失值表示了这些损失的综合。

loss_D_B:0.xxxxxx 这是判别器B（Discriminator B）的损失。判别器B用于判断B类图片的真伪，它的损失同样由多个部分组成。这里的损失值表示了这些损失的综合。

在训练过程中，这些损失值会随着训练的进行而变化。理想情况下，希望生成器的损失逐渐减小，而判别器的损失逐渐增大，这表明生成器在生成逼真的图像，而判别器在正确地判断真实和虚假的图像。

请注意，这些损失值是经过归一化处理的，因此它们的具体数值并不是绝对的，而是相对于训练数据的。在不同的训练过程中，这些损失值会有所不同。