Создание 3D-сетки из изображения с помощью Python

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

Введение

Монокулярная оценка глубины — это задача оценки значения глубины (расстояния относительно камеры) каждого пикселя для одного (монокулярного) изображения RGB. Результат такой оценки — карта глубины, которая в основном представляет собой матрицу, где каждый элемент соответствует спрогнозированной глубине соответствующего пикселя входного изображения:

Карта глубины

Точки на карте глубины можно рассматривать как множество точек с координатами по трём осям. Карта — это матрица, а значит, каждый её элемент имеет компоненты x, y (столбец и строка соответственно), а z — значение спрогнозированной глубины в точке (x, y). Список точек (x, y, z) в области обработки трёхмерных данных называется облаком точек.

Облако точек. Оригинальный файл Open3D

Можно начать с неструктурированного облака точек и получить сетку, то есть трёхмерное представление объекта из множества вершин и многоугольников [полигонов]. Самый распространённый тип сетки — треугольная сетка, состоящая из множества соединённых общими рёбрами или вершинами трёхмерных треугольников. В литературе вы найдёте несколько методов получения треугольной сетки из облака точек; самые популярные — альфа-форма¹, шаровое вращение² и реконструкция поверхности Пуассона³. Эти методы называют алгоритмами реконструкции поверхности.

Треугольная сетка. Оригинальный файл Open3d

Процедура создания сетки из изображения в этом руководстве состоит из трёх этапов:

1. Оценка глубины: с использованием монокулярной модели оценки глубины создаётся карта глубины входного изображения.
2. Построение облака точек: карта глубины преобразуется в облако точек.
3. Генерация сетки: с помощью алгоритма реконструкции поверхности из облака точек создаётся сетка.

Чтобы выполнить эту процедуру, вам понадобится изображение. Если его нет под рукой, скачайте его здесь:

Спальня. Изображение из NYU-Depth V2

1. Оценка глубины

Выбранная для этого руководства модель монокулярной оценки глубины — GLPN⁴. Получить её можно в Hugging Face Model Hub с помощью библиотеки Transformers от Hugging Face.

Для этого установите последнюю версию Transformers из PyPI:

pip install transformers

Приведённый ниже код оценивает глубину входного изображения:

import matplotlib
matplotlib.use('TkAgg')
from matplotlib import pyplot as plt
from PIL import Image
import torch
from transformers import GLPNFeatureExtractor, GLPNForDepthEstimation

feature_extractor = GLPNFeatureExtractor.from_pretrained("vinvino02/glpn-nyu")
model = GLPNForDepthEstimation.from_pretrained("vinvino02/glpn-nyu")

# load and resize the input image
image = Image.open("image.jpg")
new_height = 480 if image.height > 480 else image.height
new_height -= (new_height % 32)
new_width = int(new_height * image.width / image.height)
diff = new_width % 32
new_width = new_width - diff if diff < 16 else new_width + 32 - diff
new_size = (new_width, new_height)
image = image.resize(new_size)

# prepare image for the model
inputs = feature_extractor(images=image, return_tensors="pt")

# get the prediction from the model
with torch.no_grad():
    outputs = model(**inputs)
    predicted_depth = outputs.predicted_depth

# remove borders
pad = 16
output = predicted_depth.squeeze().cpu().numpy() * 1000.0
output = output[pad:-pad, pad:-pad]
image = image.crop((pad, pad, image.width - pad, image.height - pad))

# visualize the prediction
fig, ax = plt.subplots(1, 2)
ax[0].imshow(image)
ax[0].tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
ax[1].imshow(output, cmap='plasma')
ax[1].tick_params(left=False, bottom=False, labelleft=False, labelbottom=False)
plt.tight_layout()
plt.pause(5)

Для работы с GLPN библиотека Transformers предоставляет два класса: GLPNFeatureExtractor — для предварительной обработки входных данных, и класс модели — GLPNForDepthEstimation.

Из-за архитектуры выходной размер модели составляет:

Выходной размер. Изображение сгенерировано с помощью CodeCogs

Таким образом, размер image изменяется так, чтобы высота и ширина были кратны 32, а иначе выходные данные модели окажутся меньше входных. Это необходимо, поскольку облако точек будет прорисовано с использованием пикселей изображения, для чего входное изображение и выходная карта глубины должны обладать одинаковым размером.

Монокулярные модели оценки глубины пытаются получить прогнозы высокого качества вблизи границ, поэтому выходные данные (output) обрезаются по центру (строка 33). Чтобы сохранить одинаковые размеры, также обрезается по центру image (строка 34).

Вот некоторые прогнозы:

Прогноз глубины спальни. На входе изображение из NYU-Depth V2

Прогноз глубины игровой комнаты. На входе изображение из NYU-Depth V2

Прогноз глубины офиса. На входе изображение из NYU-Depth V2

2. Построение облака точек

В части 3D-обработки будет использоваться Open3d⁵. Наверное, это лучшая библиотека Python для задач такого рода.

Установите последнюю версию Open3d из PyPI:

pip install open3d

Код ниже преобразует предполагаемую карту глубины в объект облака точек Open3D:

import numpy as np
import open3d as o3d

width, height = image.size

depth_image = (output * 255 / np.max(output)).astype('uint8')
image = np.array(image)

# create rgbd image
depth_o3d = o3d.geometry.Image(depth_image)
image_o3d = o3d.geometry.Image(image)
rgbd_image = o3d.geometry.RGBDImage.create_from_color_and_depth(image_o3d, depth_o3d, convert_rgb_to_intensity=False)

# camera settings
camera_intrinsic = o3d.camera.PinholeCameraIntrinsic()
camera_intrinsic.set_intrinsics(width, height, 500, 500, width/2, height/2)

# create point cloud
pcd = o3d.geometry.PointCloud.create_from_rgbd_image(rgbd_image, camera_intrinsic)

Изображение RGBD — это просто комбинация RGB-изображения и соответствующего изображения глубины. Класс PinholeCameraIntrinsic хранит так называемую внутреннюю матрицу камеры. С этой матрицей Open3D может создать облако точек из изображения RGBD с правильным расстоянием между точками. Внутренние параметры оставьте как есть. Дополнительные сведения смотрите в дополнительных ресурсах в конце руководства.

Для визуализации выполните эту строку:

o3d.visualization.draw_geometries([pcd])

3. Генерация сетки

Среди различных методов для этой задачи, которые вы найдёте в литературе, здесь используется алгоритм реконструкции поверхности Пуассона³: он обычно даёт результаты лучше и мягче других.

С помощью алгоритма из полученного на последнем шаге облака точек Пуассона этот код генерирует сетку:

# outliers removal
cl, ind = pcd.remove_statistical_outlier(nb_neighbors=20, std_ratio=20.0)
pcd = pcd.select_by_index(ind)

# estimate normals
pcd.estimate_normals()
pcd.orient_normals_to_align_with_direction()

# surface reconstruction
mesh = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pcd, depth=10, n_threads=1)[0]

# rotate the mesh
rotation = mesh.get_rotation_matrix_from_xyz((np.pi, 0, 0))
mesh.rotate(rotation, center=(0, 0, 0))

# save the mesh
o3d.io.write_triangle_mesh(f'./mesh.obj', mesh)

Сначала код удаляет выбросы из облака точек. Облако может содержать шум и артефакты по разным причинам. В этом сценарии модель могла бы предсказать некоторые глубины, которые слишком сильно различаются с соседними глубинами.

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

Наконец, алгоритм выполняется. Уровень детализации сетки определяется значением depth. Помимо повышения качества сетки более высокое значение глубины увеличивает размеры вывода.

Для визуализации сетки советую скачать MeshLab, потому что есть программы 3D-визуализации только в ч/б.

Вот окончательный результат:

Сгенерированная сетка

Сетка с другого ракурса

Поскольку окончательный результат изменяется в зависимости от значения depth, это сравнение его различных значений:

Сравнение различных значений глубины

Алгоритм с depth=5 привёл к сетке в 375 КБ, depth=6 — к 1,2 МБ, depth=7 — к 5 МБ, depth=8 — к 19 МБ, depth=9 — к 70, а depth=10 — к 86 МБ.

Вывод

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

Дополнительные ресурсы:

• Модель камеры-обскуры;
• Внутренние и внешние матрицы;
• Нормали;
• Официальная документация Open3D.

Спасибо, что прочитали. Надеюсь, вы нашли материал полезным.

А мы научим работать с Python, чтобы вы прокачали карьеру или стали востребованным IT-специалистом:

• Профессия Fullstack-разработчик на Python (16 месяцев)
• Профессия Data Scientist (24 месяца)

Чтобы посмотреть все курсы, кликните по баннеру: