Пристальный взгляд на код из лучшего доклада конференции по компьютерному зрению и распознаванию образов 2021 — GIRAFFE

25 июня завершилась конференция CVPR – 2021, и какая замечательная подборка докладов! Глубокое обучение продолжает доминировать в области компьютерного зрения: у нас есть новые методы для SLAM, оценки позы, оценки глубины, новые наборы данных, сети GAN, а также многочисленные доработки прошлогодних нейронных полей свечения^[1] — NeRF, и это далеко не всё.

Возможно, вы уже слышали о работе GIRAFFE^[2]. Получив главный приз за лучшую работу этого года, она объединяет сети GAN, NeRF и дифференцируемый рендеринг, чтобы генерировать новые изображения. Однако, что важнее, новый подход предоставляет модульный фреймворк конструирования и композиции трёхмерных сцен из объектов в полностью дифференцируемом и обучаемом стиле — и это на шаг приближает нас к миру нейронного 3D-дизайна. К старту курса о машинном и глубоком обучении делимся переводом статьи, автор которой подробно рассматривает исходный код GIRAFFE и создаёт несколько кратких примеров визуализаций. На КДПВ вы видите кадр из презентации GIRAFFE.

Нейронные поля свечения

Наглядное объяснение и демонстрация NeRF

Говоря коротко, NeRF представляет собой метод описания и визуализации трёхмерной сцены в терминах её плотности и излучения в любой заданной точке трёхмерного объёма. Она тесно связана с концепцией световых полей, то есть функций, выражающих, как свет проходит через данное пространство.

Для заданной точки (x,y,z) в пространстве изобразим луч с направлением (θ, φ) на сцену. Для каждой точки вдоль луча соберём её плотность и зависимое от вида излучаемое свечение в этой точке, затем объединим эти лучи в одно пиксельное значение, как при обычной трассировке лучей. При этом сцены NeRF обучаются на коллекции снятых в различных позах изображений объектов, похожих на те, что применяются в приложениях типа «структура исходя из движения».

GIRAFFE

Наглядное объяснение и демонстрация GIRAFFE

Обзор 

В сущности, GIRAFFE — это основанный на обучении, полностью дифференцируемый механизм рендеринга, позволяющий составлять сцену как совокупность нескольких "полей признаков", то есть обобщение полей сияния в NeRFs. Эти поля признаков представляют собой трёхмерные объёмы, где каждый воксел содержит вектор признака.

Поля признаков строятся путём композиции созданных GAN обученных представлений, принимающих латентные коды в качестве входа в трёхмерную сцену. Поля признаков применяются к 3D-объёму, поэтому можно применять преобразования подобия, такие как поворот, перенос и масштабирование. Можно даже составить целую сцену как совокупность отдельных полей характеристик. По сравнению с NeRF этот метод даёт такие преимущества:

• Может представлять несколько объектов и один фон с независимыми преобразованиями (оригинальный NeRF поддерживает только одну «сцену» и не отделяет объекты друг от друга).

• Может применять позы и преобразования подобия — поворот, перенос и масштабирование — к отдельным объектам.

• Создающие поля признаков сети GAN могут независимо обучаться и повторно использоваться в качестве компонентов.

• Имеет дифференцированный движок рендеринга со сквозным обучением.

• Значения цвета не ограничиваются RGB и могут распространяться на другие свойства материала.

• Для кодирования положения использует позиционное кодирование, как в трансформере, что также “вводит индуктивное смещение для изучения представлений трёхмерных форм в канонических ориентациях, которые иначе оказались бы произвольными”.

Проект GIRAFFE содержит исходный код, он может использоваться для воспроизведения фигур из проекта и даже для создания ваших сцен. Я дам краткое руководство по их исходному коду и покажу, как работать с GIRAFFE — создавать простые нейронные трёхмерные сцены.

Исходный код

Репозиторий GIRAFFE структурирован с учётом конфигурации. Файл configs/default.yaml определяет конфигурации приложения по умолчанию. Другие файлы конфигурации, например configs/256res/cars_256/pretrained.yaml, наследуют содержимое от этого файла при помощи ключа inherit_from и переопределяют значения по умолчанию, указывая другие пары «ключ — значение».

Этот подход позволяет не составлять входные параметры отдельно, а вместо этого выводить изображения, запустив скрипт с параметром render.py <CONFIG.yaml> и обучать сети запуском скрипта с параметром — train.py <CONFIG.yaml>.

Чтобы посмотреть на рендеры в деле, сначала выполните инструкции по быстрому запуску в файле README.md: так вы загрузите предварительно обученную модель и запишете ряд выходных визуализаций, они показаны ниже.

Файл конфигурации просто берёт значения по умолчанию и вставляет предварительно обученную на наборе данных Cars модель. Этот файл создаёт довольно много визуализаций различных манипуляций с рендерингом, среди них — интерполяция внешнего вида, интерполяция формы, интерполяция фона, вращение и перенос. Эти визуализации задаются в файле configs/default.yaml ключом render_program, значение которого — список определяющих визуализации строк. Они определяют «программы рендеринга», которые рендер GIRAFFE будет вызывать, обращаясь к render.py.

В методе render_full_visualization метода im2scene.giraffe.rendering.Renderer вы увидите ряд операторов if, которые ищут имена ещё большего количества программ рендеринга: object_translation_circle, render_camera_elevation и render_add_cars. Давайте посмотрим на них в деле. Создадим новый файл конфигурации с именем cars_256_pretrained_more.yaml и добавим в него такие строки:

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
inherit_from: configs/256res/cars_256.yaml
training:
  out_dir:  out/cars256_pretrained
test:
  model_file: https://s3.eu-central-1.amazonaws.com/avg-projects/giraffe/models/checkpoint_cars256-d9ea5e11.pt
rendering:
  render_dir: rendering
  render_program: ['render_camera_elevation', 'render_add_cars']

Это предыдущий файл, он работал с ключом render_program нашего стандартного файла; мы просто перезаписали в него новые программы рендеринга. Теперь, чтобы получить больше визуализаций, выполним такую команду:

python render.py configs/256res/cars_256_pretrained_more.yaml

Должно получиться что-то вроде этого:

И вот так:

Как эти программы рендеринга на самом деле размещают, переносят и поворачивают эти автомобили? Чтобы ответить на этот вопрос, внимательнее посмотрим на класс Renderer. В примере с рендером object_rotation выше вызывается метод Renderer.render_object_rotation.

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
class Renderer(object):
    # ...
    def render_object_rotation(self, img_out_path, batch_size=15, n_steps=32):
        gen = self.generator
        bbox_generator = gen.bounding_box_generator

        n_boxes = bbox_generator.n_boxes

        # Set rotation range
        is_full_rotation = (bbox_generator.rotation_range[0] == 0
                            and bbox_generator.rotation_range[1] == 1)
        n_steps = int(n_steps * 2) if is_full_rotation else n_steps
        r_scale = [0., 1.] if is_full_rotation else [0.1, 0.9]

        # Get Random codes and bg rotation
        latent_codes = gen.get_latent_codes(batch_size, tmp=self.sample_tmp)
        bg_rotation = gen.get_random_bg_rotation(batch_size)

        # Set Camera
        camera_matrices = gen.get_camera(batch_size=batch_size)
        s_val = [[0, 0, 0] for i in range(n_boxes)]
        t_val = [[0.5, 0.5, 0.5] for i in range(n_boxes)]
        r_val = [0. for i in range(n_boxes)]
        s, t, _ = gen.get_transformations(s_val, t_val, r_val, batch_size)

        out = []
        for step in range(n_steps):
            # Get rotation for this step
            r = [step * 1.0 / (n_steps - 1) for i in range(n_boxes)]
            r = [r_scale[0] + ri * (r_scale[1] - r_scale[0]) for ri in r]
            r = gen.get_rotation(r, batch_size)

            # define full transformation and evaluate model
            transformations = [s, t, r]
            with torch.no_grad():
                out_i = gen(batch_size, latent_codes, camera_matrices,
                            transformations, bg_rotation, mode='val')
            out.append(out_i.cpu())
        out = torch.stack(out)
        out_folder = join(img_out_path, 'rotation_object')
        makedirs(out_folder, exist_ok=True)
        self.save_video_and_images(
            out, out_folder, name='rotation_object',
            is_full_rotation=is_full_rotation,
            add_reverse=(not is_full_rotation))
    # ...

Этот метод генерирует диапазон матриц вращения r для членов заданного пакета. Затем он итеративно передаёт члены этого диапазона (и некоторые значения по умолчанию для масштабирования и переноса) в метод сети GAN — forward, который задаётся ключом generator в файле default.yaml. Если теперь вы посмотрите на im2scene.giraffe.models.__init__.py, то увидите, что этот ключ сопоставлен с im2scene.giraffe.models.generator.Generator.

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
from im2scene.giraffe.models import generator
# ...
generator_dict = {
    'simple': generator.Generator,
}

Потерпите меня, пока смотрите на Generator.forward. Он принимает различные необязательные входные аргументы, такие как transformations, bg_rotation и camera_matrices, а затем передаёт их в свой метод volume_render_image, где происходит магия композиции. Латентные коды всех объектов сцены, включая наш фон, разбиваются на составляющие их формы и внешнего вида.

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
z_shape_obj, z_app_obj, z_shape_bg, z_app_bg = latent_codes

Здесь латентный код генерируется случайным образом при помощи функции torch.randn:

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
class Generator(nn.Module):
    # ...
    def get_latent_codes(self, batch_size=32, tmp=1.):
        z_dim, z_dim_bg = self.z_dim, self.z_dim_bg
        n_boxes = self.get_n_boxes()
        def sample_z(x): return self.sample_z(x, tmp=tmp)
        z_shape_obj = sample_z((batch_size, n_boxes, z_dim))
        z_app_obj = sample_z((batch_size, n_boxes, z_dim))
        z_shape_bg = sample_z((batch_size, z_dim_bg))
        z_app_bg = sample_z((batch_size, z_dim_bg))
        return z_shape_obj, z_app_obj, z_shape_bg, z_app_bg

    def sample_z(self, size, to_device=True, tmp=1.):
        z = torch.randn(*size) * tmp
        if to_device:
            z = z.to(self.device)
        return z
    # ...

А здесь прямой проход декодера отображает точки трёхмерного пространства и направление обзора камеры в значения σ и RGB (признак) для каждого объекта. К фону применяется другой генератор (для удобства детали опущены).

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
n_iter = n_boxes if not_render_background else n_boxes + 1
# ...
for i in range(n_iter):
    if i < n_boxes:  # Object
        p_i, r_i = self.get_evaluation_points(pixels_world,
            camera_world, di, transformations, i)
        z_shape_i, z_app_i = z_shape_obj[:, i], z_app_obj[:, i]
        feat_i, sigma_i = self.decoder(p_i, r_i, z_shape_i, z_app_i)
        # ...
    else:  # Background
        p_bg, r_bg = self.get_evaluation_points_bg(pixels_world,
            camera_world, di, bg_rotation)
        feat_i, sigma_i = self.background_generator(
            p_bg, r_bg, z_shape_bg, z_app_bg)
        # ...
    feat.append(feat_i)
    sigma.append(sigma_i)
# ...

Затем, с помощью σ max либо среднего значения при помощи функции composite_function из этих отображений составляется композиция.

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
sigma_sum, feat_weighted = self.composite_function(sigma, feat)

Окончательное изображение создаётся путём взвешивания отображения признаков по объёму σ  вдоль вектора луча. Результат — один кадр в одном окне анимаций выше. Чтобы больше узнать о том, как построены di и ray_vector, смотрите generator.py,

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
weights = self.calc_volume_weights(di, ray_vector, sigma_sum)
feat_map = torch.sum(weights.unsqueeze(-1) * feat_weighted, dim=-2)

Подводя итоги, давайте попробуем создать собственную программу рендеринга: чтобы добиться эффекта вращения и скольжения автомобиля слева направо, просто комбинируем вращение и перенос глубины. Для этого напишем несколько простых дополнений к классу Renderer в rendering.py.

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
class Renderer(object):
    # ...
    def render_full_visualization(self, img_out_path,
            render_program=['object_rotation']):
        for rp in render_program:
            # ...
            # APPEND THIS TO THE END OF render_full_visualization
            if rp == 'object_wipeout':
                self.set_random_seed()
                self.render_object_wipeout(img_out_path)
    # ...
    # APPEND THIS TO THE END OF rendering.py
    def render_object_wipeout(self, img_out_path, batch_size=15,
            n_steps=32):
        gen = self.generator

        # Get values
        latent_codes = gen.get_latent_codes(batch_size, tmp=self.sample_tmp)
        bg_rotation = gen.get_random_bg_rotation(batch_size)
        camera_matrices = gen.get_camera(batch_size=batch_size)
        n_boxes = gen.bounding_box_generator.n_boxes
        s = [[0., 0., 0.]
             for i in range(n_boxes)]
        n_steps = int(n_steps * 2)
        r_scale = [0., 1.]

        if n_boxes == 1:
            t = []
            x_val = 0.5
        elif n_boxes == 2:
            t = [[0.5, 0.5, 0.]]
            x_val = 1.0

        out = []
        for step in range(n_steps):
            # translation
            i = step * 1.0 / (n_steps - 1)
            ti = t + [[0.1, i, 0.]]
            # rotation
            r = [step * 1.0 / (n_steps - 1) for i in range(n_boxes)]
            r = [r_scale[0] + ri * (r_scale[1] - r_scale[0]) for ri in r]

            transformations = gen.get_transformations(s, ti, r, batch_size)
            with torch.no_grad():
                out_i = gen(batch_size, latent_codes, camera_matrices,
                            transformations, bg_rotation, mode='val')
            out.append(out_i.cpu())
        out = torch.stack(out)

        out_folder = join(img_out_path, 'object_wipeout')
        makedirs(out_folder, exist_ok=True)
        self.save_video_and_images(
            out, out_folder, name='object_wipeout',
            add_reverse=True)

Скопируйте эти дополнения в rendering.py и создайте файл конфигурации configs/256res/cars_256_pretrained_wipeout.yaml:

# adapted from https://github.com/autonomousvision/giraffe (MIT License)
inherit_from: configs/256res/cars_256.yaml
training:
  out_dir:  out/cars256_pretrained
test:
  model_file: https://s3.eu-central-1.amazonaws.com/avg-projects/giraffe/models/checkpoint_cars256-d9ea5e11.pt
rendering:
  render_dir: rendering
  render_program: ['object_wipeout']

Выполнив python render.py configs/256res/cars_256_pretrained_wipeout.yaml, вы должны получить примерно такой результат:

Заключение

GIRAFFE — это захватывающее пополнение в потоке последних исследований в области NeRF и GAN. Представление поля сияния описывает мощный, расширяемый фреймворк, с помощью которого возможно строить трёхмерные сцены в обучаемом и дифференцируемом стиле.

Статья открывает интересные возможности для экспериментов с изображениями и напоминает о том, насколько быстро развивается глубокое обучение. Но чтобы двигать область вперёд по-прежнему нужны люди. Если вам инетесна сфера машинного и глубокого обучения, то вы можете присмотреться к программе курса «Machine Learning и Deep Learning», где рассматривается множество различных нейронных сетей, включая сети GAN, а если вас интересует лаконичный Python, вы можете обратить внимание на наш курс о Fullstack-разработке на этом языке.

Узнайте, как прокачаться и в других специальностях или освоить их с нуля:

• Профессия Data Scientist

• Профессия Data Analyst

• Курс по Data Engineering