查看原文
其他

使用条件GAN实现图像到图像的翻译

gloomyfish OpenCV学堂 2020-02-04

点击上方蓝字关注我们

星标或者置顶【OpenCV学堂】

干货与教程第一时间送达!

概述

前面写过两篇关于GAN的文章,一直觉得GAN是视觉领域十分有用的利器,前一阵子看到一个通过手绘图纸找实物工业零件的公司,我再次被图像翻译这个领域的创新落地震惊啦。作为图像翻译应用经典模型pix2pix,刚出来就读了论文,但是一直没写出点什么,也没有自己跑过,这次我重新读了一遍,有跑了一下tensorflow版本的代码实现,只能用“神奇”两个字来形容。不信看下图

图像处理、视觉领域的很多问题都可以看成是翻译问题,就像把一种语言翻译成另外一种语言一样。比如灰度图像彩色化、航空图像区域分割、设计图的真实虚拟等,跟语言翻译一样,很少有一对一的直接翻译。图像整合了梯度信息、边缘信息、色彩与纹理信息,传统的图像翻译基于像素级别无法有效建模,而条件生成对抗网络(Conditional GANs)可以对这类问题有很好的效果。

基本思想


GAN中的生成者是一种通过随机噪声学习生成目标图像的模型,而条件GAN主要是在生成模型是从观察到的图像与随机噪声同时学习生成目标图像的模型,生成者G训练生成输出图像尝试让它与真实图像无法被鉴别者D区分、而鉴别者D训练学习如何区分图像是真实的还是来自生成者G。

条件GAN可以表达为:

G尝试最小化生成损失、生成目标图像、而D尝试最大化鉴别图像是否来自生成者G,对比正常的GAN表达为

此外在损失计算当中,还发现对比标注数据、加入L2或者L1的损失可以让效果更好,而且L1的效果比L2好,L2总体会让图像更加的模糊。

网络架构


无论是生成者还是鉴别者,都采用卷积网络的

CONV + BN + ReLU

形式实现网络模型拼接。


其中生成者有两种典型的结构

基于UNet的结构保留了输入信息采用skip-connection的策略进行合并,因而效果更好。

 

鉴别者网络的设计过程中,作者借鉴了马尔可夫随机场理论,认为只有相邻的像素块/像素之间有相互关系,鉴别者不再基于整张图像进行,而是基于NxN的像素快(Patch)该方法又称为Patch GAN,运行得到每个Patch得到响应之后取平均值作为D的最终输出。根据选择N大小不一样可以分为像素级别、像素块级别、图像级别等不同Patch的鉴别者。它们结构如下(D:表示鉴别者):

70x70 D
C64-C128-C256-C512

1x1 D(像素级别鉴别者)
C64-C128 – (1x1卷积)

16x16 D
C64-C128

256x256 D
C64-C128-C256-C512-C512-C512

所有ReLU都是leaky的,


不同的Patch最终生成的图像效果不一样!

代码实现

生成器G:

class Generator(tf.keras.Model):
    def __init__(self):
        super(Generator, self).__init__()
        initializer = tf.random_normal_initializer(0., 0.02)

        self.down1 = Downsample(644, apply_batchnorm=False)
        self.down2 = Downsample(1284)
        self.down3 = Downsample(2564)
        self.down4 = Downsample(5124)
        self.down5 = Downsample(5124)
        self.down6 = Downsample(5124)
        self.down7 = Downsample(5124)
        self.down8 = Downsample(5124)

        self.up1 = Upsample(5124, apply_dropout=True)
        self.up2 = Upsample(5124, apply_dropout=True)
        self.up3 = Upsample(5124, apply_dropout=True)
        self.up4 = Upsample(5124)
        self.up5 = Upsample(2564)
        self.up6 = Upsample(1284)
        self.up7 = Upsample(644)

        self.last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS,
                                                    (44),
                                                    strides=2,
                                                    padding='same',
                                                    kernel_initializer=initializer)

    @tf.contrib.eager.defun
    def call(self, x, training):
        # x shape == (bs, 256, 256, 3)
        x1 = self.down1(x, training=training)  # (bs, 128, 128, 64)
        x2 = self.down2(x1, training=training)  # (bs, 64, 64, 128)
        x3 = self.down3(x2, training=training)  # (bs, 32, 32, 256)
        x4 = self.down4(x3, training=training)  # (bs, 16, 16, 512)
        x5 = self.down5(x4, training=training)  # (bs, 8, 8, 512)
        x6 = self.down6(x5, training=training)  # (bs, 4, 4, 512)
        x7 = self.down7(x6, training=training)  # (bs, 2, 2, 512)
        x8 = self.down8(x7, training=training)  # (bs, 1, 1, 512)

        x9 = self.up1(x8, x7, training=training)  # (bs, 2, 2, 1024)
        x10 = self.up2(x9, x6, training=training)  # (bs, 4, 4, 1024)
        x11 = self.up3(x10, x5, training=training)  # (bs, 8, 8, 1024)
        x12 = self.up4(x11, x4, training=training)  # (bs, 16, 16, 1024)
        x13 = self.up5(x12, x3, training=training)  # (bs, 32, 32, 512)
        x14 = self.up6(x13, x2, training=training)  # (bs, 64, 64, 256)
        x15 = self.up7(x14, x1, training=training)  # (bs, 128, 128, 128)

        x16 = self.last(x15)  # (bs, 256, 256, 3)
        x16 = tf.nn.tanh(x16)

        return x16


判别器D:

class Discriminator(tf.keras.Model):
    def __init__(self):
        super(Discriminator, self).__init__()
        initializer = tf.random_normal_initializer(0., 0.02)

        self.down1 = DiscDownsample(644, False)
        self.down2 = DiscDownsample(1284)
        self.down3 = DiscDownsample(2564)

        # we are zero padding here with 1 because we need our shape to
        # go from (batch_size, 32, 32, 256) to (batch_size, 31, 31, 512)
        self.zero_pad1 = tf.keras.layers.ZeroPadding2D()
        self.conv = tf.keras.layers.Conv2D(512,
                                           (44),
                                           strides=1,
                                           kernel_initializer=initializer,
                                           use_bias=False)
        self.batchnorm1 = tf.keras.layers.BatchNormalization()

        # shape change from (batch_size, 31, 31, 512) to (batch_size, 30, 30, 1)
        self.zero_pad2 = tf.keras.layers.ZeroPadding2D()
        self.last = tf.keras.layers.Conv2D(1,
                                           (44),
                                           strides=1,
                                           kernel_initializer=initializer)

    @tf.contrib.eager.defun
    def call(self, inp, tar, training):
        # concatenating the input and the target
        x = tf.concat([inp, tar], axis=-1)  # (bs, 256, 256, channels*2)
        x = self.down1(x, training=training)  # (bs, 128, 128, 64)
        x = self.down2(x, training=training)  # (bs, 64, 64, 128)
        x = self.down3(x, training=training)  # (bs, 32, 32, 256)

        x = self.zero_pad1(x)  # (bs, 34, 34, 256)
        x = self.conv(x)  # (bs, 31, 31, 512)
        x = self.batchnorm1(x, training=training)
        x = tf.nn.leaky_relu(x)

        x = self.zero_pad2(x)  # (bs, 33, 33, 512)
        # don't add a sigmoid activation here since
        # the loss function expects raw logits.
        x = self.last(x)  # (bs, 30, 30, 1)

        return x


构建UNet网络作为生成者G的时候卷积与转置卷积层代码实现如下:

class Downsample(tf.keras.Model):
    def __init__(self, filters, size, apply_batchnorm=True):
        super(Downsample, self).__init__()
        self.apply_batchnorm = apply_batchnorm
        initializer = tf.random_normal_initializer(0., 0.02)

        self.conv1 = tf.keras.layers.Conv2D(filters,
                                            (size, size),
                                            strides=2,
                                            padding='same',
                                            kernel_initializer=initializer,
                                            use_bias=False)
        if self.apply_batchnorm:
            self.batchnorm = tf.keras.layers.BatchNormalization()

    def call(self, x, training):
        x = self.conv1(x)
        if self.apply_batchnorm:
            x = self.batchnorm(x, training=training)
        x = tf.nn.leaky_relu(x)
        return x


class Upsample(tf.keras.Model):
    def __init__(self, filters, size, apply_dropout=False):
        super(Upsample, self).__init__()
        self.apply_dropout = apply_dropout
        initializer = tf.random_normal_initializer(0., 0.02)

        self.up_conv = tf.keras.layers.Conv2DTranspose(filters,
                                                       (size, size),
                                                       strides=2,
                                                       padding='same',
                                                       kernel_initializer=initializer,
                                                       use_bias=False)
        self.batchnorm = tf.keras.layers.BatchNormalization()
        if self.apply_dropout:
            self.dropout = tf.keras.layers.Dropout(0.5)

    def call(self, x1, x2, training):
        x = self.up_conv(x1)
        x = self.batchnorm(x, training=training)
        if self.apply_dropout:
            x = self.dropout(x, training=training)
        x = tf.nn.relu(x)
        x = tf.concat([x, x2], axis=-1)
        return x


判别器与生成器的损失计算

def discriminator_loss(disc_real_output, disc_generated_output):
  real_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels = tf.ones_like(disc_real_output),
                                              logits = disc_real_output)
  generated_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels = tf.zeros_like(disc_generated_output),
                                                   logits = disc_generated_output)
  total_disc_loss = real_loss + generated_loss
  return total_disc_loss


def generator_loss(disc_generated_output, gen_output, target):
  gan_loss = tf.losses.sigmoid_cross_entropy(multi_class_labels = tf.ones_like(disc_generated_output),
                                             logits = disc_generated_output)
  # mean absolute error
  l1_loss = tf.reduce_mean(tf.abs(target - gen_output))
  total_gen_loss = gan_loss + (LAMBDA * l1_loss)
  return total_gen_loss


训练与运行演示

因为我的机器比较慢(GTX1050ti),基于CMP_facade数据集运行40个epoch的时候输出如下:


作者论文的模型一些测试效果如下,真的是太cool了,不信请看下图。

是不是想自己运行一下啦,代码我已经提交到github上啦,这里下载即可

https://github.com/gloomyfish1998/dl_learning_notes


几乎是各种图像风格的完美翻译与实现,是真正的pix2pix的神奇模型!最后如果觉得文章不错,请点【在看】支持



往期精选



参考资料:

论文: https://arxiv.org/abs/1611.07004
源码: https://github.com/phillipi/pix2pix


关注【OpenCV学堂】

长按或者扫码即可关注

    您可能也对以下帖子感兴趣

    文章有问题?点此查看未经处理的缓存