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

В этой заметке мы рассмотрим какие существуют основные датасеты в области 3D ML и какие фреймворки для работы с 3D данными могут пригодиться датасаентисту при разработке моделей машинного обучения в данной области.

Серия 3D ML на Хабре:

1. Формы представления 3D данных
2. Функции потерь в 3D ML
3. Датасеты и фреймворки в 3D ML

Репозиторий на GitHub для данной серии заметок.

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

Datasets

В зависимости от типа входных и выходных данных, у алгоритма машинного обучения могут существенно разниться способы решения задач, подсчета метрик качества и организации датасета. В связи с этим важно понимать особенности устройства датасета, который будет использован при обучении модели в конкретной задаче. Драфт таксономии задач 3D ML представлен на рис.1. Больше информации об этих задачах можно найти на этой github странице или в обзорных дакладах с тематических конференций (сайт geometricdeeplearning.com раздел Tutorials).

Рис.1 Таксономия задач 3D ML.

С момента появления в сети первых трехмерных моделей, которые исследователи по всему миру стали использовать для своих экспериментов (Стэнфордский кролик, Чайник из Юты), датасеты трехмерных моделей претерпели значительное развитие. Во-первых, это касается размеров датасетов: если в первых датасетах количество моделей исчислялось сотнями, то сейчас речь идет о миллионах моделей (как в ShapeNet). Во-вторых, в современных датасетах помимо самих моделей, хранящихся в какой либо форме представления данных, находится различная дополнительная информация, позволяющая решать более сложные типы задач.

Рис.2 Слайд о развитии 3D датасетов из доклада Михаила Бронштейна о Geometric Deep Learning на SGP 2018.

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

Здесь содержится представительная выборка датасетов с описанием и ссылками для задач 3D ML, также в данном репозитории можно найти много другой полезной информации, например, ссылки на материалы для изучения, примеры архитектур и всевозможных проектов.

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

1. ShapeNet (2015) [1]

Датасет содержит более 3 миллионов моделей (.obj формат), в более чем 4 тысячах категорий. На странице проекта есть удобный браузер для предпросмотра моделей. Датасет хорошо организован и содержит много дополнительной информации о моделях: текстуры, описания, снимки моделей. Используется для проведения соревнований (SHREC, ICCV) по задачам 3D ML.

В свободном распространении доступны более компактные версии датасета:

• ShapeNetCore: содержит 51300 моделей в 55 категориях.
• ShapeNetSem: содержит 12000 более аннотированных моделей в 270 категориях.

Недостатком данного датасета в контексте задачи 2D-to-3D является отсутствие фотографий реальных объектов, соответствующих 3D моделям. Также тот факт, что модели были созданы художниками, а не получены в результате сканирования, может повлиять на тестирование моделей на реальных данных.

2. ModelNet (2015) [2]

Датасет содержит 127915 3D моделей в 662 категориях (в основном файлы в формате .off).

Имеются более компактные версии датасета:

• ModelNet10: 4899 моделей в 10 категориях.
• ModelNet40: 12311 моделей в 40 категориях, все модели одинаково ориентированы.

На странице проекта можно ознакомиться с таксономией моделей в датасете, найти всю сопутствующую документацию и ознакомиться с результатами соревнования “ModelNet Benchmark”. В датасете так же как и в ShapeNet есть изображения только отрендеренных моделей, но нет изображений реальных объектов, соответствующих этим моделям.

3. Pix3D (2018) [3]

Датасет содержит 395 3D моделей в .obj формате и вокселизованные версии моделей в .mat формате (.mat является Matlab форматом, но его можно свободно прочитать с помощью scipy.io). Модели представлены в 9 категориях и для каждой модели имеется набор реальных изображений в высоком качестве, с различным контекстом окружающего пространства. Помимо изображений и меток классов, в датасете присутствует информация о ключевых точках объектов как для изображений, так и для 3D моделей.

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

4. ABC dataset (2019) [4]

Датасет содержит порядка миллиона CAD моделей в параметрическом виде. Каждая модель состоит из набора аналитических примитивов, таких как плоскость, эллипс, цилиндр, тор и пр.

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

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

Рис.3 Сравнение 3D датасетов по качественным характеристикам и объему [4].

5. VOCASET: Speech-4D Head Scan Dataset (2019) [5]

VOCASET — 4D face dataset с примерно 29-минутным сканированием лиц разговаривающих людей, снятым со скоростью 60 кадров в секунду и синхронизированным звуком. Набор данных содержит 12 субъектов и 480 последовательностей продолжительностью около 3-4 секунд каждая с предложениями, выбранными из массива стандартных протоколов, которые максимизируют фонетическое разнообразие.

6. ScanNet (2017) [6]

Набор данных состоящий из RGB-D видео, содержащий 2,5 миллиона кадров для более чем 1500 отсканированных объектов, аннотированных позициями 3D-камер, реконструкциями поверхностей и семантическими сегментациями сцен.

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

7. Semantic3D (2017) [7]

Крупномасштабный бенчмарк для классификации облаков точек, который предоставляет большой аннотированный набор 3D данных в виде сцен, представленных в виде облака точек улиц и внешних экстерьеров (outdoor scene) сцен с более 4 миллиардов точек, а также охватывает целый ряд разнообразных городских сцен.

Наравне с датасетом KITTY, Semantic3D может использоваться в задачах, связанных с беспилотным транспортом.

8. Campus3D (2020) [8]

Данный датасет на сегодняшний день является самым большим датасетом экстерьеров в форме облака точек. Исследователи из Национального Университета Сингапура с помощью аэрофотосъемки и технологии фотограмметрии смогли перенести весь университетский кампус в аннотированную 3D модель облака точек (с более чем 900 миллионами точек).

Рис.4 Сравнение популярных датасетов в форме облака точек из статьи [8].

Frameworks

Рис.5 Часть модулей, составляющих библиотеку Kaolin [9].

После того, как исследователь определился с задачей, архитектурой, метриками, функциями потерь и датасетом, ему необходимо выстроить удобный и эффективный процесс обучения и тестирования (т.н. train/test pipelines). Удобнее всего пользоваться готовыми фреймворками, приспособленными для работы в конкретной предметной области. Для задач классического компьютерного зрения некоторые наиболее распространенные фреймворки для обучения моделей приведены в нашей заметке об опыте использования машинного зрения в промышленной дефектоскопии.

Конечно, существует много различного проприетарного и свободно распространяемого ПО, которое может быть полезным для предобработки, визуализации и даже генерации 3D контента. Из свободно распространяемых и наиболее употребимых отметим Blender, PCL (Point cloud library) и MeshLab. Ниже, речь пойдет в основном о библиотеках для Python, которые являются дополнениями к существующим тензорным фреймворкам (в основном Tensorflow и Pytorch) библиотеках для решения вспомогательных задач.

Отдельно стоит упомянуть роль игровых движков и в частности Unity 3D с ее библиотекой Unity ML agents в развитии задач машинного обучения, связанных с пространством (например RL задачи про шагающие механизмы и многое другое). Подробнее про применение ML agents в RL можно прочитать здесь: раз, два, три, четыре.

Для нетерпеливых более компактной формы восприятия, все рассматриваемые далее фреймворки с их основными характеристиками мы собрали в одну таблицу.

Фреймворки основные и вспомогательные

Библиотека, которая выполняет роль основной при решении задач 3D ML зачастую содержит следующие основные части (см. пример модулей из Kaolin на рис.5):

• Загрузчик данных для наиболее распространенных форматов аннотации 3D датасетов;
• Визуализатор данных (может быть как интерактивным так и не интерактивным);
• Дифференциальный рендеринг — система рендеринга трехмерных сцен, позволяющая строить непрерывные дифференцируемые зависимости между параметрами основных объектов сцены (освещение, текстуры, камеры и меш объектов) и их растрового отображения. Дифференциальный рендер позволяет создавать модели нейронного рендеринга — системы рендеринга, в которых отдельные части заменены нейронными сетями, что позволяет выполнять задачу отрисовки трехмерных сцен быстрее и разнообразнее чем в классических системах рендеринга. (больше деталей про дифференциальный рендеринг можно посмотреть здесь: видео 1, видео 2, видео 3; примеры из библиотек: redner, softras, pytorch3d expl; medium paper);
• Расчет метрик и функций потерь для 3D данных;
• Библиотека архитектур (model zoo) машинного обучения являющихся SOTA для своих задач в области работы с 3D данными;
• Поддержка различных форм представления 3D данных и способов конвертации их друг в друга;
• Дополнительные утилиты для выполнения специфических задач.

Рис.6 Схема прямого и обратного рендеринга внутри фреймворка. Иллюстрация из доклада об устройстве PyTorch3D.

Чтобы лучше понять взаимосвязь описанных выше модулей и то как они устроены в базовых фреймворках, можно ознакомиться со статьей [9] про NVIDIA Kaolin, или со статьей [10] и видеопрезентацией про фреймворк PyTorch3D, или с видеопрезентацией про TensorFlow Graphics [11].

К вспомогательным фреймворкам стоит отнести библиотеки, которые решают задачи, которые не так сильно коррелируют именно с частью машинного обучения: это могут быть библиотеки для интерактивной визуализации, такие как polyscope, или библиотеки, которые нацелены на решение задач связанных со специфическим форматом данных, такие как mesh_to_sdf.

Основные фреймворки

1. PyTorch Geometric (Fey & Lenssen: Department of Computer Graphics TU Dortmund University | 2019) [12]

Данная библиотека — надстройка над PyTorch, приспособлена для работы с неевклидовыми структурами данных, такими как пространственные графы. В библиотеке реализованы различные операторы сверток на графах и ряд архитектур глубокого обучения, основанных на этих операторах. Присутствуют загрузчики датасетов и конструктор по созданию собственного датасета.

Несмотря на то, что PyTorch Geometric создавалась как библиотека для 3D ML, присутствующие в ней модели не охватывают весь спектр задач данной области машинного обучения. В библиотеке рассмотрена только две формы представления трехмерных данных — полигональный меш и облако точек. Отсутствует модуль для графического представления данных и дифференциального рендеринга. Библиотека может использоваться совместно с PyTorch3D.

Доступные платформы: Linux; Windows; Mac (CPU only).

Примеры с кодом использования фреймворка.

2. TensorFlow Graphics (Google Brain | 2019) [11]

Так же как и PyTorch Geometric, TensorFlow Graphics создавалась для работы с пространственными графами.

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

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

Доступные платформы: Linux; Windows; Mac.

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

Рис.6 Внешний вид TensorBoard 3D. Пример доступен в репозитории проекта.

3. NVidia Kaolin (NVidia | 2019) [9]

Хорошо продуманная библиотека для работы с задачами 3D ML, которая находиться в бета версии (крайне неустойчива и функционал реализован не полностью). В библиотеке реализованы все основные модули необходимые для работы с задачами данной предметной области — от интерактивной визуализации моделей и дифференциального рендеринга, до набора глубоких архитектур и загрузчика датасетов.

Несмотря на продуманность и разнообразие модулей в библиотеке, пока что в ней присутствуют баги, недостаточная документированность и реализована не вся заявленная функциональность.

Доступные платформы: Linux; Windows (unstable).

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

Рис.7 Архитектура дифференциального рендера Kaolin [9].

4. PyTorch 3D (Facebook research | 2019) [10]

Несмотря на устойчивую работу только под Linux и отсутствие некоторого функционала, библиотека активно развивается, хорошо приспособлена под работу на кластерах видеокарт и в совокупности со вспомогательными библиотеками (наприм. PyTorch Geometric + trimesh + polyscope) может в будущем послужить базовой библиотекой для работы в области 3D ML.

Есть поддержка как CPU так и GPU. Некоторые утилиты, были специально оптимизированы для задач 3D ML: это и быстрое вычисление chamfer loss, и нахождение ближайших соседей для точки из облака точек, и продуманный механизм дифференциального рендеринга, и операторы свертки на графах, и представление батча 3D моделей в виде PyTorch тензоров.

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

Доступные платформы: Linux; Mac; Windows (unstable).

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

5. Points 3D (Chaton & Chaulet: Principia Labs | 2020)

Библиотека предназначена для решения задач глубокого обучения применительно к облакам точек. Присутствуют SOTA архитектуры, аугментация данных и интерактивная визуализация.
Доступна на всех ОС как для CPU так и для GPU. Если работаете только с облаками точек, можно использовать в качестве основного фреймворка.

Более подробное описание библиотеки можно посмотреть в заметке на Medium.

Доступные платформы: Linux; Mac; Windows.

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

Рис.8 Расширение Jupyter Lab для интерактивной визуализации моделей в виде облаков точек, на которых можно проверить работу обученных в Points 3D архитектур.

Вспомогательные фреймворки

1. Kornia (2019) [50]

Область 3D ML можно в некотором смысле рассматривать как развитие классического машинного зрения, поэтому неудивительно, что некоторый функционал для решения задач geometrical deep learning можно найти уже в библиотеке компьютерного зрения Kornia для PyTorch.

Рис.9 Сравнение библиотек компьютерного зрения [13].

В целом, несмотря на то, что среди библиотек классического компьютерного зрения, Kornia является наиболее проработанной и функциональной, ее возможности как библиотеки 3D ML весьма ограничены, и из поддерживаемых трехмерных типов данных имеются только RGB-D изображения.

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

2. Polyscope

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

Рис.10 Демонстрация работы интерактивной визуализации меша в Polyscope.

3. trimesh

В библиотеке содержится множество функций для разнородной обработки меша:

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

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

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

4. mesh_to_sdf

Рис.11 Пример расчета SDF функции с помощью mesh_to_sdf в точках вне и внутри полигональной модели стула.

Библиотека предназначена для вычисления значений SDF функции в пространстве вокруг произвольного меша. Библиотека работает в том числе с полигональными моделями которые в англоязычной литературе принято называть non-watertight meshes: в моделях могут присутствовать дыры, некорректно ориентированные полигоны, самопересекающиеся полигоны, полигональные модели которые обладают т.н. non-manifold геометрией.

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

Пример того как использовать функционал данной библиотеки в глубоком обучении можно рассмотреть на примере проекта DeepSDF [14].

5. Open3D [15]

Аналог рассмотренного выше trimesh. Есть имплементация как для Python так и для C++. Библиотека содержит широкий функционал для работы с полигональными моделями и облаками точек.

Примеры использования фреймворка можно найти в репозитории проекта или в статье [15], посвященной фреймворку.

Выводы

Данные и способ их организации — закономерно центральные темы в deep learning и machine learning. 3D ML не исключение, поэтому важно понимать какие данные и датасеты использовать для обучение своей модели, где их взять и как предобрабатывать, визуализировать с помощью специальных фреймворков и библиотек.

Универсального инструмента, который позволяет решить все встающие перед исследователем задачи нет и даже все перечисленные выше библиотеки не всегда содержат в себе необходимый функционал, однако некоторую часть задач они все же покрывают и могут послужить хорошей основой для написание собственного решения. Из нашего собственного опыта в качестве основной библиотеки можно использовать PyTroch3D + оригинальные имплементации изучаемых архитектур, в качестве вспомогательных библиотек можно использовать trimesh, PCL и Open3D для предобработки меша и облаков точек (для более быстрой работы можно пользоваться C++ версиями библиотек), для визуализации можно пользоваться Python API у Blender, но для более простых визуализаций может быть достаточно инструментария базовых фреймворков.

Дополнительно рекомендуем ознакомиться с видео туториалом "SGP 2020 Graduate School: Black Box Geometric Computing with Python", в котором авторы АВС датасета рассказывают о библиотеках для Python, которые они разработали сами для того, чтобы работать с 3D данными при прототипировании алгоритмов машинного обучения.