2D-to-3D: конструируем сервис для экспериментов с реконструкциями формы

Привет, хабровчане! Сегодня мы хотим поделиться с вами нашим опытом по созданию сервиса для апробации моделей 3D реконструкции. В заметки мы:

• кратко обсудим что из себя представляет задача 2D-to-3D, 

• взглянем на наиболее успешные алгоритмы и работы, 

• сравним параметрический и непараметрический подходы к восстановлению формы лица и тела человека,

• разберем как устроен наш сервис для апробации моделей пространственной реконструкции.

Всех, кто хотел бы на манер “Чудес науки” переместить девушку с обложки глянцевого журнала к себе в комнату хотя бы в AR, просим под кат =).

Заметка от партнера IT-центра МАИ и организатора магистерской программы “VR/AR & AI” — компании PHYGITALISM.

Задача 2D-to-3D

Если коротко: задача 2d-to-3d в своем названии содержит суть проблемы - у нас имеются 2D данные (это могут быть одно или множество изображений объекта,  видео, на котором объект снимают с разных сторон,набросок контура объекта и другие подобные данные) и мы хотим, чтобы алгоритм машинного обучения по этому изображению или по метке класса объекта восстановил 3D геометрию в том или ином виде.

Существуют схожие задачи, в которых также восстанавливают 3D форму объектов, но на входе у алгоритмов чуть больше информации: например это может быть не просто RGB изображение, а RGB-D (то есть по разряженному облаку точек).

Из обзорных работ, вышедших за последние годы, в которых есть сравнения и обсуждения разных подходов к данной задаче, стоит выделить “What do single-view 3d reconstruction networks learn?” [1], в которой авторы предложили собственные алгоритмы для восстановления полигонального меша объекта по единичному изображению и сравнили его с другими SOTA моделями 2019 года. Кстати, задача которая будет рассмотрена далее как раз относится к такому же более узкому классу задач Single-view 3D reconstruction.

Из общих моментов в этой задаче стоит отметить, что существующие подходы основаны на:

• разных обучениях представлениям (где-то это Encoder-Decoder, где-то GAN и т.д.), 

• разных формах представления 3D (в некоторых подходах используются воксели, в некоторых меш, а в некоторых даже параметрическое представление)

• разных функциях ошибок.

Про виды обучения представлениям можно прочитать в ставшей уже классикой статье “Representation learning: A review and new perspectives” [3]. Различные формы представления 3D данных и различные функции потерь для мы освещали в наших заметках здесь (раз и два). Также стоит упомянуть что на paperswithcode существует отдельная страница, посвященная именно этой задаче.

Примеры архитектур для 3D Reconstruction

Рассмотрим для примера несколько наиболее известных архитектур 2D-to-3D. Условно, можно разделить все архитектуры по признакам:

• конечное 3D представление (меш, облако точек, воксели и пр.), 

• параметрическая или непараметрическая модель,

• модель общего назначения или приспособленная под определенные классы объектов.

1. Mesh R-CNN [2019/Facebook research] [4]

В данной работе исследователи из Facebook AI Research применяют Mask R-CNN [5] (про данную архитектуру есть много хороших обзорных материалов, на русском языке, например, смотрите вот эту заметку), чтобы выделить маски объектов на изображениях. После чего на основе этих масок генерируется воксельная сетка будущей модели, а на последнем этапе вершины вокселей трансформируются с помощью графовых сверток и модель становится полигональной. Стоит отметить разумный подход к выбору функции ошибок и регуляризаторам.

На GitHub странице проекта доступен Collab, в котором можно протестировать модель на своих изображениях. 

2. Occupancy networks [2019] [6]

В данной модели в качестве формы представления 3D объектов используется вероятность принадлежности данной точки пространства внутренности объекта:

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

Второе преимущество заключается в том, что нейронная сеть за счет малого и фиксированного числа входных параметров очень компактна.

Стоит отметить, что благодаря тому, что сеть по сути аппроксимирует некоторую вероятностную функцию, которая описывает форму и способ кодирования входной информации для этой формы - \x может быть произвольным. Такой подход позволяет решать самые разные задачи из области 3D ML, что и было продемонстрировано авторами.

3. Atlas [2020][7]

Данная работа хоть и не относится к Single-view 3D reconstruction, но является интересным примером использования другого 3D представления - TSDF (truncated signed distance field / воксельная сетка пространства с вероятностью существования точки поверхности в каждом вокселе). В данном подходе на каждом кадре  формируется проекционным методом TSDF, а потом TSDF со всех кадров агрегируются и получается итоговая TSDF вместе с метками класса для каждого вокселя на основе меток класса в 2D изображениях.

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

Восстановление геометрии человеческого тела

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

Мы уже рассматривали параметрическую модель поверхности лица человека в заметке про дифференциальный рендеринг, где использовали классическую модель, полученную в институте Макса Планка еще в 1999 г. Рассмотрим более актуальные модели, основанные на том же параметрическом принципе, который кратко можно выразить следующим образом:

"Геометрия некоторой поверхности представляет из себя параметрическое семейство поверхностей, которое для каждого фиксированного значения параметра фиксируется в допустимой конфигурации."

1. DECA (2018)[9]

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

2. HMR (2018)[8]

Параметрическая модель тела человека с оцененными ключевыми точками. Модель основана на минимизации отклонения пар ключевых точек, благодаря чему для ее обучения можно обойтись без 3D датасета. На видео выше представлена демонстрация применения модели для видео, которое мы выбрали для тестов. 

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

3. PiFU HD (Facebook AI Research/2020)[11]

Данная непараметрическая модель реконструкции  и текстуры тела человека от исследователей из Facebook AI Research является усовершенствование их предыдущей работы PiFU [10].

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

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

Конструируем свой веб сервис для 2D-to-3D моделей

Предположим, что вы исследователь, который хочет протестировать большое число моделей или разработчик, которому нужно создать простенький сервис по тестированию моделей 2D-to-3D, как вы будете это делать? Перед нами тоже встала такая задача, и мы выбрали путь конвейерной обработки данных, где каждый этап представлял из себя отдельный Docker контейнер, а для визуализации и взаимодействия с сервисом  - web клиент, написанный на JS. Давайте пройдемся по нашему 2D-to-3D сервису и посмотрим как он устроен.

Клиент взаимодействует с сервисом через браузер на своем устройстве. Фронтенд при этом написан на JS и выполнен минималистично в виде одного экрана для выбора 2D-to-3D модели и ее параметров, загрузки данных (в данном примере для модели PiFU HD это будет фотография человека). Сервер реализован на Python и крутится внутри Docker контейнера. После того, как сервер примет данные от пользователя, он сохранит их в хранилище данных и по соответствующей очереди RabbitMQ (который также крутиться в отдельном Docker контейнере) отправит задание на обработку в сервис 2D-to-3D.

Сначала (хотя это и не обязательно для всех моделей) изображение нужно обрезать и оставить только ту часть, где присутствует человек (получить кроп и выделить маску). Для этого можно в Docker контейнер обернуть любой ваш любимый фреймворк современного компьютерного зрения, мы использовали gluon-cv. Полученные после препроцессинга данные также помещаются в хранилище к исходным данным.

После предобработки изображения задание на реконструкцию поступает в алгоритм восстановления геометрии (в нашем случаи PiFU HD). Мы модифицировали исходный код авторов для того, чтобы встроить его в наш конвейер и оптимизировали память и быстродействие. Поскольку для реализации многих алгоритмов 3D ML требуются вычисления на видеокарте, данный модуль реализован в виде nvidia-docker контейнера.

После получения 3D моделей мы сохраняем их в .glb формате и потом используем возможности отображения моделей в AR в браузере (дополнительно конвертируем модели в .usdz для того, чтобы iOS устройства также могли работать с моделями, однако стоит учитывать какие при этом появляются ограничения).

В будущем можно будет легко заменить или добавить модули в такой конвейер, чтобы работать с множеством различных архитектур и делиться своими экспериментами с миром (что мы вообщем то и стараемся делать). Надеемся, что вам понравилась идея учить машины “лепить” скульптуры по изображению. Будем рады, если вы поделитесь своими идеями, где можно применить подобные модели на практике.

Источники:

1. Tatarchenko M. et al. What do single-view 3d reconstruction networks learn? //Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. – 2019. – С. 3405-3414. [paper]

2. Firman, Michael and Mac Aodha, Oisin and Julier, Simon and Brostow, Gabriel J. (2016). Structured Prediction of Unobserved Voxels From a Single Depth Image / Firman CVPR 2016 [code]

3. Bengio Y., Courville A., Vincent P. Representation learning: A review and new perspectives //IEEE transactions on pattern analysis and machine intelligence. – 2013. – Т. 35. – №. 8. – С. 1798-1828. [paper]

4. Gkioxari G., Malik J., Johnson J. Mesh r-cnn //Proceedings of the IEEE/CVF International Conference on Computer Vision. – 2019. – С. 9785-9795. [paper][code]

5. Kaiming He, Georgia Gkioxari, Piotr Dollar, Ross Girshick, “Mask R-CNN”. arXiv:1409.4842 2017-2018 [paper][code]

6. Mescheder L. et al. Occupancy networks: Learning 3d reconstruction in function space //Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. – 2019. – С. 4460-4470. [paper][code]

7. Murez Z. et al. Atlas: End-to-end 3d scene reconstruction from posed images //arXiv preprint arXiv:2003.10432. – 2020. [project page]

8. Kanazawa A. et al. End-to-end recovery of human shape and pose //Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition. – 2018. – С. 7122-7131. [project page]

9. Feng Y. et al. Learning an Animatable Detailed 3D Face Model from In-The-Wild Images //arXiv preprint arXiv:2012.04012. – 2020. [paper][code]

10. Saito S. et al. Pifu: Pixel-aligned implicit function for high-resolution clothed human digitization //Proceedings of the IEEE/CVF International Conference on Computer Vision. – 2019. – С. 2304-2314. [project page]

11. Saito S. et al. Pifuhd: Multi-level pixel-aligned implicit function for high-resolution 3d human digitization //Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. – 2020. – С. 84-93. [project page]