В компьютерном зрении часто приходится работать с двумерными изображениями, и значительно реже - с 3D объектами. Из-за этого многие ML инженеры чувствуют себя неуверенно в этой области: много незнакомых слов, непонятно, куда тут применить старых друзей Resnet и Unet. Поэтому сегодня я хотел бы немного поговорить о 3D на примере задачи определения шести степеней свободы, что в каком-то виде синонимично 3D object detection. Я разберу одну из относительно свежих работ на эту тему с некоторыми отступлениями. 

Меня зовут Арсений, я работаю ML инженером и веду Telegram-канал partially unsupervised. Эта статья написана по мотивам моего же видео для Data Fest 2020, секция “CV в индустрии”. 

Кратко о задаче

Для начала давайте определимся, что такое шесть степеней свободы (6 DoF - degrees of freedom). Представим себе некоторый ригидный (неизменяемый, т.е. при трансформации все точки будут оставаться на той же дистанции друг от друга) объект в трехмерном мире. Чтобы описать его положение относительно наблюдателя понадобится 6 измерений: три будут отвечать за повороты по разным осям, а еще три - за смещение по соответствующим осям. Соответственно, имея эти шесть чисел, мы представляем, как объект расположен относительно какого-то базиса (например, точки, с которой ведется фотосъемка). Эта задача является классической для робототехники (где находится объект, который нужно схватить роборукой?), дополненной реальности (где нарисовать маску в MSQRD, ушки в Snapchat или кроссовки в Wanna Kicks?) , беспилотных автомобилей и других доменов.

Будем рассматривать статью MobilePose: Real-Time Pose Estimation for Unseen Objects with Weak Shape Supervision (Hou et al., 2020). Эта статья, написанная авторами из Google Research, предлагает надежный и, что немаловажно, быстрый пайплайн для решения задачи, будет уместно разобрать его по частям.

Пайплайн состоит из трех основных кусков:

1. Backbone достаточно классический, архитектура в виде песочных часов должна быть знакома каждому, кто хоть раз обучал Unet. 

2. Выходы сети не выглядят инновационно. Detection, regression - все знакомые слова! Впрочем, насчет shape могут возникнуть вопросы. Но давайте отложим их на время.

3. Постпроцессинг может показаться загадочным для тех, кто не в теме. Что такое EPnP и почему оно превращает 2D точки в 3D bounding box? 

3D для чайников

И здесь сразу нужно сделать важное отступление, которое поможет нам ответить на все эти вопросы. Давайте высокоуровнево посмотрим на некоторую математику 3D мира. Пусть есть некоторый набор 3D-точек X - матрица размером (n, 3), в которой n - количество точек. Как мы помним, шесть степеней свободы - это три поворота и три смещения, т.е. Rigid transformation. Если обозначить R матрицу поворота, а t - вектор переноса (rotation и translation соответственно), будет соблюдаться такое уравнение: 

X’ = X @ R.T + t 

R и t и есть то, что мы хотим найти в этой задаче, они описывают то, как надо сдвинуть и повернуть наш ригидный объект, чтобы он оказался там, где мы его сейчас видим.

Но X’ - это все еще 3D-координаты. Потому стоит сказать, что еще существует некоторая проективная матрица P. Эта матрица характеризует то, как мы проецируем объект на двумерную плоскость, условно “рендерим” его. Эта матрица зависит от размера фотографии, фокусного расстояния, искажений, но в нашей задаче ее можно считать константной. Имя такую матрицу, можно получить 2D координаты точек, просто умножив ее на X’: 

x = X’ @ P.T

Совсем на пальцах: мы ищем, как нужно повернуть и подвинуть некий объект, чтобы какие-то его точки спроецировались так, как они сейчас изображены на фотографии. Я все упростил до неприличия, потому отправляю всех желающих просветлиться посмотреть на CS231a. 

Подзадача нахождения R и t, зная X, x и P, называется Perspective-n-Point. Т.е. мы знаем, как выглядит наш объект в 3D (это X), знаем, на какое изображение он спроецирован (P) и где на этом изображении находятся его точки. Выглядит как задача оптимизации! Есть целое семейство алгоритмов, которые решают эту задачу, например, некоторые уже реализованы в OpenCV.

Еще некоторые ссылки по теме:
Monocular Model-Based 3D Tracking of Rigid Objects: A Survey (Lepetit et. al 2005) - классический обзор;
EPnP: An Accurate O(n) Solution to the PnP Problem (Lepetit et. al 2008) - сильный бейзлайн;
PnP-Net: A hybrid Perspective-n-Point Network (Sheffer and Wiesel, 2020)? - для тех, кто хочет скрестить ужа и ежа, т.е. добавить к PnP немного диплернинга.

Кстати, к этой проблеме подходят и с другой стороны. Адепты deep learning могут найти множество статей, где используется специальный projection layer, который преобразует 2D и 3D точки друг в друга. Обычно, чтобы обучить такой слой, используют синтетические данные, т.к. 3D координаты из реального мира получать дорого и сложно. Пример такой статьи.

Где взять 3D точки?

Итак, нам нужен X - напомню, это набор 3D точек. Откуда его взять?

Самый простой вариант - пытаться найти один и тот же объект на всех изображениях. Берем какую-то 3D CAD модель (готовую, рисуем с нуля, сканируем настоящий объект специальным сканером…) и используем его (точнее, какие-то его точки) в качестве X. Иными словами, делаем явное допущение "на фотографии находится именно такой объект" - на первый взгляд, это прямо-таки нагло, но практика показывает, что для оценки 6 DoF этого достаточно.

Более сложный подход - так называемые параметризированные модели. Образцовый пример - Basel Face. Исследователи из Университета Базеля отсканировали много лиц и при помощи PCA обучили такую модель, чтобы изменение малого числа ее параметров позволяло сгенерировать эти 3D лица. Таким образом, можно крутить малое количество ручек - главных компонент и получать довольно разные модели.

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

Если наша 3D модель параметризована, ее параметры можно подбирать разными итеративными методами и выбирать такую, на которой reprojection error будет меньше. Т.е. берем некоторую модель X, решаем PnP, получаем R и t и выбираем такой X, чтобы разность x и (X @ R.T + t) @ P была минимальна, для примера можно посмотреть на procrustes analysis.

Истинные диплернеры идут дальше и в каком-то виде выучивают или 3D модель, или ее параметры. Хороший пример - известная работа DensePose от Facebook Research, которая популяризовала подход с выучиванием dense карты координат. Т.е. модель предсказывает для каждого пикселя его относительное расположение на 3D модели. Дальше можно найти соответствия и получить для каждого пикселя некоторое приближение его 3D координаты. 

В статье, которую мы изначально взялись разбирать, есть выход с таинственным названием shape. В зависимости от наличия grouth truth данных (об этом немного позже) авторы либо учат там сегментационную маску объекта (т.н. weak supervision, для улучшения сходимости), либо как раз карту координат объекта. 

Также может возникнуть вопрос - координаты каких именно точек мы хотим найти? Ответ простой: на самом деле, нам все равно. 3D модель обычно состоит из тысяч вершин, мы можем выбрать подмножество по своему вкусу. Единственный более или менее важный критерий - приблизительная равноудаленность точек друг от друга; в случае неудачной выборки решение PnP становится нестабильным. 

Где взять 2D точки?

Итак, с 3D объектом худо-бедно разобрались, давайте пойдем в более знакомую большинству CV инженеров область, т.е. вернемся в 2D и подумаем, где взять координаты на плоскости. 

Для получения 2D координат точек (обычно эта задача называется keypoint detection) популярны два основных подхода: регрессия напрямую (последний слой сети выдает x, y для каждой точки) и heatmap-регрессия (последний слой выдает тепловую карту с пятном около точки). Первый подход может быть быстрее, т.к. необязательно выстраивать полную encoder-decoder архитектуру, второй обычно точнее и достаточно легко обучается (это почти та же сегментация).

Авторы MobilePose не пошли ни по одному из этих путей и придумали интересный гибрид, вдохновившись современными “безанкорными” архитектурами для детекции вроде CenterNet. Помните, на схеме есть головы Detection и Regression? Так вот, Detection голова предсказывает центр объекта, а Regression - где вершины объекта находятся относительно этого центра. 

В этом подходе есть изящный трюк. Я много писал о том, как выбрать 3D модель, а в этом случае все осознанно упрощается: в качестве модели используется простой параллелепипед, он же 3D bounding box! То есть X - это вершины бокса, а x - проекция этих вершин. Значит, достаточно знать соотношения сторон этого бокса (которые мы получаем из самой архитектуры детектора), и остальная магия ни к чему.

Особенности приготовления данных

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

Unlike 2D object detection, it is prohibitive to manually label data for 3D detection. Due to this difficulty of collecting sufficiently large amounts of labeled training data, such approaches are typically trained on real data that are highly correlated with the test data (e.g., same camera, same object instances, similar lighting conditions). As a result, one challenge of existing approaches is generalizing to test data that are significantly different from the training set.

Synthetic data is a promising alternative for training such deep neural networks, capable of generating an almost unlimited amount of pre-labeled training data with little effort. Synthetic data comes with its own problems, however. Chief among these is the reality gap, that is, the fact that networks trained on synthetic data usually do not perform well on real data.

В рассматриваемой статье авторы сделали один из таких трюков: вместо того, чтобы бездумно рендерить весь мир, они сделали комбинацию реальных и синтетических данных. Для этого они взяли видео из AR-приложений (хорошо быть гуглом и иметь много данных из ARCore). В таких видео есть и знание о плоскостях, оценка 6 DoF, полученная при помощи визуальной одометрии, и оценка освещения. Это позволяет рендерить искусственные объекты не где попало, а только на плоских поверхностях, адаптируя освещенность, что значительно снижает reality gap между синтетическими и реальными данными. К сожалению, повторить этот трюк в домашних условиях кажется довольно сложным.

Все вместе

Ура, мы галопом пробежались по всем ключевым концепциям пайплайна! Этого должно хватить, чтобы читатель смог собрать из open source компонентов, например, приложение, которое будет рисовать маску на лице (для этого можно даже не учить модели самостоятельно, готовых сетей для face alignment немало).

Конечно, это будет только прототип. При доведении такого приложения до продакшена возникнет еще много вопросов, например:

• Как добиться консистентности между кадрами?

• Что делать с неригидными объектами?

• Что делать, если объект частично не виден?

• Что делать, если в кадре много объектов?

Именно там начнутся настоящие приключения.