Pix2Pix（从零实现图像风格转换任务）

柔光的暖阳◎ 2022-08-29 04:26 312阅读 0赞

本篇博文的内容是利用CGAN从零实现图像风格转换任务，本文用的数据集图片如下，目标是由右边的仅包含人物线条的图片生成左边的彩色图（通俗的来说就是给黑白图片上色）。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70]

各位可以根据提供的地址从Kaggle网站上将图片下载。[动漫图片数据集地址][Link 1]  
数据集下载下来有文件夹重复的情况，最后仅保留三个文件夹即可，如下图  
![在这里插入图片描述][568561677070489b99d16222bd4e268c.png]

## CGAN论文的简要介绍 ##

论文的作者提出，普通的GAN都是期望生成器能够学习到将一个噪声分布（比如高斯分布或者是均匀分布）映射成我们想要的概率分布。  
 r a n d o m   z random \\ z random z    G : z → y \\ \\bold G:z \\rarr y  G:z→y  
但是在CGAN中，我们的式子变成了  
 r a n d o m   z ,    i n p u t \_ i m a g e   x random \\ z,\\ \\ input \\\_image\\ x random z,  input\_image x    G : x , z → y \\ \\bold G:\{x, z\} \\rarr y  G:x,z→y  
这里的input\_image指的是仅包含线条的图片  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70 1]  
实作上，我参考的Github项目并没有往Generator中输入random noise z。我的理解是本来Generator的目的就是将一个分布映射成我们想要的分布，增加一个noise分布似乎并没有增添有用的信息。论文的作者也提到了增加noise以后，结果的随机性增加的有限，当然读者可以在实现时增加一个noise来对比效果。

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70 2]  
上图是Generator的结构简图，它包含了两个部分，一部分是将输入图片进行下采样的encoder，另一部分是进行上采样的decoder。作者考虑到输入（仅包含线条的图片）与输出（上了色的图片）在结构上具有对称性，比如两种图片的线条轮廓是一样的。所以作者提出了将encoder中的特征图concatenate到了decoder的对应位置上，图中的虚线就表示了这种关系。

而CGAN中discriminator的特殊之处在于，它接收的输入是 x , y x,y x,y concatenate的结果，并且它的输出不再是一个标量，而是代表一个四维的张量，论文作者将其称为一个patch。  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70 3]  
![在这里插入图片描述][72e607b078ef41a3bb1758ec576f4984.png]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70 4]  
以上就是关于模型设计的全部信息，不过在实现时，**我们将BatchNorm替换成InstanceNorm**。

### 模型设计 ###

#Generator
    import torch
    from torch import nn
    class CNNBlock(nn.Module):
    	def __init__(self, in_channels, out_channels, relu=False, down=True, use_dropout=False):
    		#in_channels输入的通道数
    		#out_channels输出的通道数
    		#relu 是用relu还是leakyReLU
    		#down 是下采样还是上采样
    		#use_dropout是否是用dropout
    		self.block = nn.Sequential(
    			nn.Conv2d(in_channels, out_channels, 4, 2, 1, bias=True)
    			if down else
    			nn.ConvTranspose2d(in_channels, out_channels, 4, 2, 1, bias=False, padding_mode='reflect'),
    			nn.InstanceNorm2d(out_channels),
    			nn.Dropout(0.5) if use_dropout,
    			nn.ReLu() if relu else nn.LeakyReLu(0.2)
    	)
    class Generator(nn.Module):
    	def __init__(self):
    		super().__init__(img_channels=3, feature=64)
    		#encoder一开始没有batchnorm
    		self.down1 = nn.Sequential(
    			nn.Conv2d(img_channels, feature, 4, 2, 1, padding_mode='reflect'),
    			nn.LeakyReLU(0.2)
    		)
    		self.down2 = CNNBLock(feature, feature * 2)
    		self.down3 = CNNBlock(feature * 2, feature * 4)
    		self.down4 = CNNBlock(feature * 4, feature * 8)
    		self.down5 = CNNBlock(feature * 8, feature * 8)	
    		self.down6 = CNNBlock(feature * 8, feature * 8)	
    		self.down7 = CNNBlock(feature * 8, feature * 8)	
    		self.bottleneck = nn.Sequential(
    			nn.Conv2d(feature * 8, feature * 8, 4, 2, 1, padding_mode='reflect'),
    			nn.ReLU()
    		)
    		#decoder部分,仅有前三层是用dropout
    		self.up1 = CNNBlock(feature * 8, feature * 8, relu=True, down=False, use_dropout=True)
    		self.up2 = CNNBlock(feature * 16, feature * 8, relu=True, down=False, use_dropout=True)
    		self.up3 = CNNBlock(feature * 16, feature * 8, relu=True, down=False, use_dropout=True)	
    		self.up4 = CNNBlock(feature * 16, feature * 8, relu=True, down=False)	
    		self.up5 = CNNBlock(feature * 16, feature * 4, relu=True, down=False)	
    		self.up6 = CNNBlock(feature * 8, feature * 2, relu=True, down=False)	
    		self.up7 = CNNBlock(feature * 4, feature, relu=True, down=False)
    		#最后一层是用Tanh作为激活函数
    		self.fianl_conv = nn.Sequential(
    			nn.ConvTranspose2d(feature * 2, img_channels, 4, 2, 1, padding_mode='reflect'),
    			nn.Tanh()
    	)
    	def forward(self, x):
    	#x[N, img_channels, 256, 256]
    		down1 = self.down1(x)#[N, 64, 128, 128]
    		down2 = self.down2(down1)#[N, 128, 64, 64]
    		down3 = self.down3(down2)#[N, 256, 32, 32]
    		down4 = self.down4(down3)#[N, 512, 16, 16]
    		down5 = self.down5(down4)#[N, 512, 8, 8]
    		down6 = self.down6(down5)#[N, 512, 4, 4]
    		down7 = self.down7(down6)#[N, 512, 2, 2]
    		bottleneck = self.bottleneck(down7)#[N, 512, 1, 1]
    		up1 = self.up1(bottleneck) #[N, 512, 2, 2]
    		up2 = self.up2(torch.cat([up1, down7], 1))#[N, 512, 4, 4]
    		up3 = self.up3(torch.cat([up2, down6], 1))#[N, 512, 8, 8]
    		up4 = self.up4(torch.cat([up3, down5], 1))#[N, 512, 16, 16]
    		up5 = self.up5(torch.cat([up4, down4], 1))#[N, 256, 32, 32]
    		up6 = self.up6(torch.cat([up5, down3], 1))#[N, 128, 64, 64]
    		up7 = self.up7(torch.cat([up6, down2], 1))#[N, 64, 128, 128]
    		final_conv = self.final_conv(torch.cat([up7, down1]))#[N, 3, 128, 128]
    		return final_conv

Disciminator就是一个很简单的下采样

import torch
    from torch import nn
    class CNNBlock(nn.Module):
    	def __init(self, in_channels, out_channels, stride):
    		super().__init__()
    		self.conv = nn.Sequential(
    			nn.Conv2d(in_channels, out_channels, 4, 2, stride, bias=False, padding_mode='reflect')
    			nn.InstanceNorm2d(out_channels),
    			nn.LeakyReLU(0.2)
    	)
    class Discriminator(nn.Module):
    	def __init__(self, img_channles, features=[64, 128, 256, 512]):
    		super().__init__()
    		self.init_conv = nn.Sequential(
    			nn.Conv2d(img_channels * 2, features[0]),
    			nn.LeakyReLU(0.2)
    	)
    	   layers = []
    	   in_channels = features[0]
    	   for feature in features:
    	   		layers.append(CNNBlock(in_channels, feature, stride=1 if feature==features[-1] else 2))
    	   		in_channels = feature
    	   self.model = nn.Sequential(*layers)
    	
    	def forward(self, x, y):
    		x = torch.cat([x, y], 1)
    		x = self.init_conv(x)
    		return self.model(x)

### 准备数据集 ###

当我们的数据集下载好以后，需要处理成Dataset类，方便DataLoader加载

#config.py
    #这里会进行一些数据增强的操作
    import albumentations as A
    from albumentations.pytorch import ToTensorV2
    #两种图片都会进行的操作
    both_transforms = A.Compose([
    			#调整大小，50%进行水平旋转
    			A.Resize(width=256, height=256), A.HorizontalFlip(p=0.5)，
    		]， additional_targets = { "image0":"image"})
    transform_only_input = A.Compose([
    	A.ColorJitter(p=0.1),
    	A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.),
    	ToTensorV2()
    ])
    transform_only_mask = A.Compose([
    	A.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], max_pixel_value=255.),
    	ToTensorV2()
    ])

import torch
    from torch.utisl.data import Dataset
    import config
    import numpy as np
    import os
    from PIL import Image
    class AnimeDataset(Dataset):
    	def __init__(self, root_dir):
    		self.root_dir = root_dir
    		self.paths = os.listdir(root_dir)
    	def __len__(self):
    		return len(self.paths)
    	def __getitem__(self, index):
    		img_path = os.path.join(self.root_dir, self.paths[index])
    		img = np.array(Image.open(img_path))
    		#图片总宽1024，右边512是input，左边512是目标 
    		input_image = img[:, 512:, :]
    		target_image = img[:, :512, :]
    		augmentations = config.both_transforms(image=input_image, image0=target_image)
    		input_image, target_image = augmentations['image'], augmentations['image0']
    		input_image = config.transform_only_input(image=input_image)['image']
    		target_image = config.transform_only_mask(image=target_image)['image']
    		return input_image, target_image

### 一些训练的时候将会用到的工具 ###

**config就是一个配置文件，记录一些超参数还有数据增强的代码**

#config.py
    import torch
    import albumentations as A
    from albumentations.pytorch import ToTensorV2
    DEVICE = "cuda" if torch.cuda.is_available() else 'cpu'
    
    LEARNING_RATE = 2e-4
    BATCH_SIZE = 16
    NUM_WORKERS = 2
    IMAGE_SIZE = 256
    CHANNELS_IMG = 3
    L1_LAMBDA = 100
    NUM_EPOCHS = 500
    LOAD_MODEL = True
    SAVE_MODEL = True
    CHECKPOINT_DISC = "disc.pth.tar"
    CHECKPOINT_GEN = "gen.pth.tar"
    both_transforms = A.Compose(
        [A.Resize(width=IMAGE_SIZE, height=IMAGE_SIZE), A.HorizontalFlip(p=0.5)], additional_targets={ "image0":"image"}
    )
    transform_only_input = A.Compose(
        [
            A.ColorJitter(p=0.1),
            A.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5, 0.5, 0.5], max_pixel_value=255.0),
            ToTensorV2()
        ]
    )
    transform_only_mask = A.Compose(
        [
            A.Normalize(mean=[0.5, 0.5, 0.5],std=[0.5, 0.5, 0.5], max_pixel_value=255.0),
            ToTensorV2()
        ]
    )

一些在训练中可能用到的工具

#utils.py
    import torch
    import config
    from torchvision.image import save_image
    
    def save_some_examples(gen, val_loader, epoch, folder):
    	x, y = next(iter(val_loader))
    	x, y = x.to(config.DEVICE), y.to(config.DEVICE)
    	gen.val()
    	with torch.no_grad():
    		y_fake = gen(x)
    		y_fake = y_fake * 0.5 + 0.5
    		save_image(y_fake, folder + f"/y_gen_{ epoch}.png")
            save_image(x * 0.5 + 0.5, folder + f"/input_{ epoch}.png")
       	    if epoch == 1:
            	save_image(y * 0.5 + 0.5, folder + f"/label_{ epoch}.png")
            	gen.train()
    
    def save_checkpoint(model, optimizer, filename):
    	checkpoint = { 
    	"model":model.state_dict(),
    	"optimizer":optimizer.state_dict()
    	}
    	torch.save(checkpoint, filename)
    
    def load_checkpoint(model, optimizer, filename, lr):
    	state_dict = torch.load(filename, map_location=config.DEVICE)
    	model.load_state_dict(state_dict['model'])
    	optimizer.load_state_dict(state_dict['optimizer'])
    	for param_group in optimizer.param_group():
    		param_group['lr'] = lr

### 训练过程 ###

import torch
    from torch import optim, nn
    from generator import Generator
    from discriminator import Discriminator
    from dataset import AnimeDataset
    from torch.utils.data import DataLoader
    from utils import *
    import config
    from tqdm import tqdm
    
    disc = Discriminator(img_channels=conifg.CHANNELS_IMAGE).to(config.DEVICE)
    gen = Generator(img_channels=config.CHANNELS_IMAGE).to(config.DEVICE)
    opt_disc = optim.Adam(disc.parameters(), lr=config.LEARNING_RATE, betas=(0.5,0.999))
    opt_gen = optim.Adam(gen.parameters(), lr=config.LEARNING_RATE, betas=(0.5,0.999))
    train_set = AnimeDataset(root_dir='../data/anime/train')
    train_loader = DataLoader(train_set, batch_size=config.BATCH_SIZE, shuffle=True, num_workers=config.NUN_WORKERS)
    val_set = AnimeDataset(root_dir='../data/anime/val')
    val_loader = DataLoader(val_set, batch_size=1, shuffle=False)
    
    bce = nn.BCEWithLogitsLoss() #损失函数依旧是js散度
    L1_loss = nn.L1Loss() #l1正则项
    if config.LOAD_MODEL:
    	load_checkpoint(gen, opt_gen, config.CHECKPOINT_GEN, config.LEARNING_RATE)
    	load_checkpoint(disc, opt_disc, config.CHECKPOINT_GEN, config.LEARNING_RATE)
    
    for epoch in range(1, config.NUM_EPOCHS + 1):
    	loop = tqdm(train_loader ,leave=True)#显示训练进度条
    	for x, y in loop:
    		x ,y = x.to(config.DEVICE), y.to(config.DEVICE)
    		y_fake = gen(x)
    		disc_real = disc(x ,y).reshape(-1)
    		disc_fake = disc(x, y_fake.detach()).reshape(-1)
    		lossD_real = bce(disc_real, torch.ones_like(disc_real))
    		lossD_fake = bce(disc_fake, torch.zeros_like(disc_fake))
    		lossD = lossD_real + lossD_fake
    		opt_disc.zero_grad()
    		lossD.backward()
    		opt_disc.step()
    		
            D_fake = disc(x, y_fake).reshape(-1)
            G_fake_loss = bce(D_fake, torch.ones_like(D_fake))
            L1 = l1(y_fake, y) * config.L1_LAMBDA
            G_loss = G_fake_loss + L1
            opt_gen.zero_grad()
            G_loss.backward()
    		opt_gen.step()
        if config.SAVE_MODEL and epoch % 5 == 0:
            save_checkpoint(gen, opt_gen, filename=config.CHECKPOINT_GEN)
            save_checkpoint(disc, opt_disc, filename=config.CHECKPOINT_DISC)
        	save_some_examples(gen, val_loader, epoch, folder='evaluation')

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70]: /images/20220829/50cc5a5563254defaca689cae185f89c.png
[Link 1]: https://www.kaggle.com/ktaebum/anime-sketch-colorization-pair
[568561677070489b99d16222bd4e268c.png]: /images/20220829/f244c60c81a640c99586cef58c2b36fb.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70 1]: /images/20220829/e29b0a99cfd242c2afcefdb02519a783.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70 2]: /images/20220829/b9132c10befd4a7fbd8a1fc88a757704.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70 3]: /images/20220829/1e164d7ccf41483ba80213d5d8d317aa.png
[72e607b078ef41a3bb1758ec576f4984.png]: /images/20220829/456a9d986c13493d8b1896417ef9268e.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3FxXzQzMTUyNjIy_size_16_color_FFFFFF_t_70 4]: /images/20220829/3c60853f7cf545e591cf496d4952a61f.png