Вы когда-нибудь пытались решить проблему перевода изображения (преобразование изображения из одной категории в другую) в своем проекте компьютерного зрения?
Если да, то вы, должно быть, поняли, что с текущими неконтролируемыми методами перевода изображения в изображение , из-за требования иметь доступ к большому количеству изображений из целевой категории (категории, в которую вы хотите перевести исходные изображения), это считается узким местом. Так как же нам подойти к решению этой проблемы?

В этой статье мы рассмотрим подход, предложенный в статье Неконтролируемая трансляция изображения в изображение за несколько кадров, опубликованной исследователями из NVIDIA, и попытаемся понять их реализацию через их код.

Эта статья является первой частью серии статей, которые я собираюсь написать в этой статье. В этой статье основное внимание уделяется объяснению части реализации «Генератор изображений».

Перевод изображений без присмотра за несколькими кадрами

Платформа FUNIT нацелена на сопоставление изображения исходного класса (скажем, кошки) с аналогичным изображением невидимого целевого класса (скажем, собаки) путем использования нескольких изображений целевого класса, доступных во время тестирования.

Предлагаемая в статье структура состоит из генератора условных изображений G и многозадачного состязательного дискриминатора D. Генератор G (называемый транслятором изображений с несколькими кадрами) одновременно принимает изображение контента «x ” и набор изображений класса K” {y1, …, yK} в качестве входных данных и создать выходное изображение “z”.

Пусть S и T обозначают набор исходных классов и набор целевых классов соответственно. Во время обучения G учится переводить изображения между двумя случайно выбранными исходными классами cx, cy ∈ S, где cx != cy. Обратите внимание, что на изображении выше мы видим, что, хотя вид источника один и тот же (собаки), они по-прежнему считаются разными классами из-за различий в их характеристиках, например, лабрадор считается другой класс, чем у хаски.
Во время тестирования G берет несколько изображений из невидимого целевого класса c ∈ T (скажем, тигров или кошек и т. д.) в качестве изображений класса и сопоставляет изображение, выбранное из любого из исходных классов, с аналогичным изображением. образ целевого класса c.

Имея это в виду, давайте попробуем понять код этой реализации.

Переводчик изображений с несколькими кадрами

Преобразователь изображений G с несколькими кадрами состоит из кодировщика контента Ex, кодировщика класса Ey и декодера Fx.
Кодер контента Ex сопоставляет входное изображение контента x со скрытым кодом контента zx, который представляет собой карту пространственных объектов.
Кодер класса Ey сначала сопоставляет каждое из K отдельных изображений класса {y1, …, yK} с промежуточным скрытым вектором, а затем вычисляет среднее значение промежуточных скрытых векторов для получения окончательного скрытый код класса zy.
Декодер Fx состоит из нескольких остаточных блоков адаптивной нормализации экземпляра (AdaIN), за которыми следует пара сверточных слоев масштабирования. Для каждого образца AdaIN сначала нормализует активации образца в каждом канале, чтобы иметь нулевое среднее значение и единичную дисперсию. Затем он масштабирует активации, используя изученное аффинное преобразование, состоящее из набора скаляров и смещений.

Код взят из официального репозитория статьи GitHub с небольшими изменениями для лучшего понимания читателей.

class FewShotGen(nn.Module):
    def __init__(self, hp):
        super(FewShotGen, self).__init__()
        nf = 64
        nf_mlp = 256
        down_class = 4
        down_content = 3
        n_mlp_blks = 3
        n_res_blks = 2
        latent_dim = 64
        self.enc_class_model = ClassModelEncoder(down_class, 
                                                 3,
                                                 nf,
                                                 latent_dim,
                                                 norm='none',
                                                 activ='relu',
                                                 pad_type='reflect')

        self.enc_content = ContentEncoder(down_content,
                                          n_res_blks,
                                          3,
                                          nf,
                                          'in',
                                          activ='relu',
                                          pad_type='reflect')

        self.dec = Decoder(down_content,
                           n_res_blks,
                           self.enc_content.output_dim,
                           3,
                           res_norm='adain',
                           activ='relu',
                           pad_type='reflect')

        # Code for MLP can be found on the github repo of the project
        # In networks.py ------> class MLP
        self.mlp = MLP(latent_dim,
                       get_num_adain_params(self.dec),
                       nf_mlp,
                       n_mlp_blks,
                       norm='none',
                       activ='relu')

    def forward(self, one_image, model_set):
        # reconstruct an image
        content, model_codes = self.encode(one_image, model_set)
        model_code = torch.mean(model_codes, dim=0).unsqueeze(0)
        images_trans = self.decode(content, model_code)
        return images_trans

    def encode(self, one_image, model_set):
        # extract content code from the input image
        content = self.enc_content(one_image)
        # extract model code from the images in the model set
        class_codes = self.enc_class_model(model_set)
        class_code = torch.mean(class_codes, dim=0).unsqueeze(0)
        return content, class_code

    def decode(self, content, model_code):
        # decode content and style codes to an image
        adain_params = self.mlp(model_code)
        assign_adain_params(adain_params, self.dec)
        images = self.dec(content)
        return images

Как упоминалось выше, генератор нескольких выстрелов состоит из кодировщика контента, кодировщика классов и декодера. MLP здесь представляет собой многоуровневый персептрон, который полагается на базовую нейронную сеть для выполнения операций.

Прежде чем разбираться в коде для кодировщиков и декодеров, давайте сначала определим отдельные блоки, которые будут использоваться для построения этих слоев.

Conv2dBlock

class Conv2dBlock(nn.Module):
    def __init__(self, in_dim, out_dim, ks, st, padding=0,
                 norm='none', activation='relu', pad_type='zero',
                 use_bias=True, activation_first=False):
        
        super(Conv2dBlock, self).__init__()
        self.use_bias = use_bias
        self.activation_first = activation_first
        # initialize padding
        if pad_type == 'reflect':
            self.pad = nn.ReflectionPad2d(padding)
        elif pad_type == 'replicate':
            self.pad = nn.ReplicationPad2d(padding)
        elif pad_type == 'zero':
            self.pad = nn.ZeroPad2d(padding)
        else:
            assert 0, "Unsupported padding type: {}".format(pad_type)

        # initialize normalization
        norm_dim = out_dim

        if norm == 'in':
            self.norm = nn.InstanceNorm2d(norm_dim)
        elif norm == 'adain':
            self.norm = AdaptiveInstanceNorm2d(norm_dim) 
        elif norm == 'none':
            self.norm = None
        else:
            assert 0, "Unsupported normalization: {}".format(norm)

        # initialize activation
        if activation == 'relu':
            self.activation = nn.ReLU(inplace=True)
        elif activation == 'tanh':
            self.activation = nn.Tanh()
        else:
            assert 0, "Unsupported activation: {}".format(activation)

        self.conv = nn.Conv2d(in_dim, out_dim, ks, st, bias=self.use_bias)

    def forward(self, x):
        x = self.conv(self.pad(x))
        if self.norm:
            x = self.norm(x)
        if self.activation:
            x = self.activation(x)
        return x

Ввод проходит через слой Conv2d, а затем проходит через нормализацию и активацию, в зависимости от переданных параметров (например, relu в качестве активации и InstanceNorm2d в качестве нормализации).

ResBlocks

class ResBlock(nn.Module):
    def __init__(self, dim, norm='in', activation='relu', pad_type='zero'):
        super(ResBlock, self).__init__()
        model = []
        model += [Conv2dBlock(dim, dim, 3, 1, 1,
                              norm=norm,
                              activation=activation,
                              pad_type=pad_type)]
        model += [Conv2dBlock(dim, dim, 3, 1, 1,
                              norm=norm,
                              activation='none',
                              pad_type=pad_type)]
        self.model = nn.Sequential(*model)

    def forward(self, x):
        residual = x
        out = self.model(x)
        out += residual
        return out

class ResBlocks(nn.Module):
    def __init__(self, num_blocks, dim, norm, activation, pad_type):
        super(ResBlocks, self).__init__()
        self.model = []
        for i in range(num_blocks):
            self.model += [ResBlock(dim,
                                    norm=norm,
                                    activation=activation,
                                    pad_type=pad_type)]
        self.model = nn.Sequential(*self.model)

    def forward(self, x):
        return self.model(x)

ResBlocks — это набор остаточных блоков (ввод + вывод операции свертки, короче) с 2 операциями свертки. Объяснение всего кода остаточного блока выходит за рамки этой статьи, но для дальнейшего чтения вы можете прочитать ЭТУ статью после того, как вы сначала закончите с этой :)

Адаптивный блок нормы экземпляра

class AdaptiveInstanceNorm2d(nn.Module):
    def __init__(self, num_features, eps=1e-5, momentum=0.1):
        super(AdaptiveInstanceNorm2d, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self.weight = None
        self.bias = None
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))

    def forward(self, x):
        b, c = x.size(0), x.size(1)
        running_mean = self.running_mean.repeat(b)
        running_var = self.running_var.repeat(b)
        x_reshaped = x.contiguous().view(1, b * c, *x.size()[2:])
        out = F.batch_norm(
            x_reshaped, running_mean, running_var, self.weight, self.bias,
            True, self.momentum, self.eps)
        return out.view(b, c, *x.size()[2:])

    def __repr__(self):
        return self.__class__.__name__ + '(' + str(self.num_features) + ')'

Чтобы узнать больше об Adaptive Instance Normalization, обратитесь к этой исследовательской статье.

Теперь, когда все приготовления сделаны, давайте начнем с основной части кода.

Кодировщик контента

class ContentEncoder(nn.Module):
    def __init__(self, downs, n_res, input_dim, dim, norm, activ, pad_type):
        super(ContentEncoder, self).__init__()
        self.model = []
        self.model += [Conv2dBlock(input_dim, dim, 7, 1, 3,
                                   norm=norm,
                                   activation=activ,
                                   pad_type=pad_type)]
        for i in range(downs):
            self.model += [Conv2dBlock(dim, 2 * dim, 4, 2, 1,
                                       norm=norm,
                                       activation=activ,
                                       pad_type=pad_type)]
            dim *= 2
        self.model += [ResBlocks(n_res, dim,
                                 norm=norm,
                                 activation=activ,
                                 pad_type=pad_type)]
        self.model = nn.Sequential(*self.model)
        self.output_dim = dim

    def forward(self, x):
        return self.model(x)

Кодировщик содержимого нацелен на извлечение скрытого инвариантного представления класса (например, позы объекта). Он состоит из нескольких ConvBlocks и Resblocks (описанных выше). Термин «контент» здесь относится к изображению контента (изображение в красном квадрате на рисунке 1), которое мы передаем в сеть во время обучения. Изображение содержимого преобразуется в карту пространственных объектов, а затем передается в декодер вместе с картой объектов кодировщика классов.

Параметры, передаваемые в блоках, являются параметрами по умолчанию из официального репозитория. Как упоминалось выше, каждый блок сам по себе является отдельным блоком (Conv2dBlock, ResBlocks), что можно увидеть в приведенном выше коде.

Кодировщик класса

class ClassModelEncoder(nn.Module):
    def __init__(self, downs, ind_im, dim, latent_dim, norm, activ, pad_type):
        super(ClassModelEncoder, self).__init__()
        self.model = []
        self.model += [Conv2dBlock(ind_im, dim, 7, 1, 3,
                                   norm=norm,
                                   activation=activ,
                                   pad_type=pad_type)]
        for i in range(2):
            self.model += [Conv2dBlock(dim, 2 * dim, 4, 2, 1,
                                       norm=norm,
                                       activation=activ,
                                       pad_type=pad_type)]
            dim *= 2
        for i in range(downs - 2):
            self.model += [Conv2dBlock(dim, dim, 4, 2, 1,
                                       norm=norm,
                                       activation=activ,
                                       pad_type=pad_type)]
        self.model += [nn.AdaptiveAvgPool2d(1)]
        self.model += [nn.Conv2d(dim, latent_dim, 1, 1, 0)]
        self.model = nn.Sequential(*self.model)
        self.output_dim = dim

    def forward(self, x):
        return self.model(x)

Во время обучения кодировщик класса учится извлекать специфичное для класса скрытое представление из изображений исходных классов. Во время тестирования это обобщается на изображения ранее невидимых классов. Кодер класса извлекает специфическое для класса скрытое представление (например, внешний вид объекта).

Декодер

class Decoder(nn.Module):
    def __init__(self, ups, n_res, dim, out_dim, res_norm, activ, pad_type):
        super(Decoder, self).__init__()

        self.model = []
        self.model += [ResBlocks(n_res, dim, res_norm,
                                 activ, pad_type=pad_type)]
        for i in range(ups):
            self.model += [nn.Upsample(scale_factor=2),
                           Conv2dBlock(dim, dim // 2, 5, 1, 2,
                                       norm='in',
                                       activation=activ,
                                       pad_type=pad_type)]
            dim //= 2
        self.model += [Conv2dBlock(dim, out_dim, 7, 1, 3,
                                   norm='none',
                                   activation='tanh',
                                   pad_type=pad_type)]
        self.model = nn.Sequential(*self.model)

    def forward(self, x):
        return self.model(x)

Это часть архитектуры декодера. Передавая скрытый код класса в декодер через слои AdaIN, мы позволяем изображениям класса управлять глобальным внешним видом (например, внешним видом объекта), в то время как изображение содержимого определяет локальную структуру (например, расположение глаз).

Конечным результатом декодера будет преобразованное изображение из исходного класса в целевой класс, имеющее «стиль» целевого класса. Итак, допустим, вы начали с изображения лабрадора в качестве исходного, в результате этой операции перевода вы получите изображение хаски (целевой класс).

Это знаменует окончание части 1 этой серии статей, часть 2 этой статьи включает в себя раздел о дискриминаторе, создании набора данных и обучении модели.
Пока вы здесь, не забудьте подписаться на меня ЗДЕСЬ на носителе, будет постоянно публиковать статьи о новых научных работах по компьютерному зрению и ИИ в целом, и если у вас есть какие-либо сомнения относительно машинного обучения или глубокого обучения, оставьте заметку здесь, я буду рад помочь вам.

Спасибо!!