KPConv на русском: Свертки для point cloud / Хабр

Приветствую всех читателей Хабра!

В одной из моих предыдущих публикаций я в конце выразил предположение о возникновении новых, более сложных задач, которые потребуют прямой работы с облаками точек, минуя преобразование их в воксели. Совсем недавно такая задача и возникла. Мне нужно было выполнить совмещение двух облаков точек (Point Cloud Registration). Задача кажется простой благодаря существованию алгоритма ICP (Iterative Closest Point), о котором я уже писал ранее, и его различных модификациях. Однако в моем случае оба облака точек имели очень малое пересечение, что значительно усложняло задачу. К тому же, на областях пересечения облако точек становилось более разряженным, что добавило еще больше трудностей.

В ходе поиска решения я наткнулся на статью под названием "RRGA-Net: Robust Point Cloud Registration Based on Graph Convolutional Attention", которая показалась мне очень интересной и по изображениям результатов хотелось попробовать на моей задаче. К сожалению, на GitHub не было доступной реализации, поэтому я решил самостоятельно имплементировать модель, описанную в статье. В процессе изучения архитектуры этой модели мне встретился незнакомый слой — KPConv, или Kernel Point Convolution.

Именно это и стало началом изучения KPConv — новым (относительно. В мире NN и за год могут происходит значительные скачки, а тут уже около 4 лет прошло) слоем нейронной сети, который я решил изучить. На русскоязычных ресурсах и на Хабре информации о KPConv практически нет, поэтому я решил взяться за задачу заполнить этот пробел на сколько это возможно в моих силах. В этой статье я опирался на оригинальные исследования, но старался изложить информацию своими словами так, чтобы она была понятна как можно широкому кругу читателей. Формулы и технические детали взяты из первоисточника, чтобы сохранить точность и ценность материала.

Основные представления 3D объектов

Начнем чуть-чуть издалека и кратко вспомним как могут быть представлены 3D данные.

Трехмерные данные могут быть представлены несколькими способами:

Облако точек (Point Cloud)

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

Воксели (Voxels)

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

Meshes

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

Что-бы понять всю ценность KPConv следует сначала кратко рассмотреть проблемы, которые появляются при работе с Point Cloud.

Трудности возникающие при работе с point cloud

Облака точек представляют собой один из самых сложных типов данных для анализа (на мой взгляд), и вот почему:

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

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

Тесно связанной с неструктурированностью является неупорядоченность данных. Облака точек не имеют естественного порядка, в котором они должны быть анализированы — каждая точка потенциально является "началом" облака. Это ставит под сомнение применение многих методов ма��инного обучения, которые зависят от последовательного или фиксированного порядка входных данных.

Основы KPConv

KPConv, или Kernel Point Convolution, представляет собой подход к свертке, который работает непосредственно с облаками точек, не требуя их предварительного преобразования в другие форматы, такие как воксели. В этом разделе мы рассмотрим основные концепции KPConv, чтобы подготовить почву для более глубокого понимания.

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

Слева логика работы свертки в изображениях.Справа логика работы свертки в Point Cloud — Слева логика работы свертки в изображениях.
Справа логика работы свертки в Point Cloud

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

Математика лежащая в основе

Рассмотрим математику, которая лежит в основе и потом покажем на примере расчет. Пример игрушечный и примитивный, но для понимания работы будет вполне достаточно .

Первым делом нам надо помнить, что положения "соседей "считается относительно той точки признаки которой мы хотим извлечь (центральная точка). Всё как и при свертке в 2D.
Позицию каждого соседа смещаем так, чтобы центральная точка оказалась в начале координат, что задается как:

Отлично! Мы узнали положение соседних точек, относительно центральной. Эти значения мы можем передать в функцию . Функция ядра присваивает различный "вес" матрицам веса ядра. Значение ядра для данного соседа определяется как взвешенная сумма всех весов ядра с соответствующими корреляциями:

$g(y_i) = \sum_{k<K} h(y_i, \tilde{x}_k) W_k$

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

Для функции корреляции используют линейную модель:

$h(y_i, \tilde{x}_k) = \max\left(0, 1 - \frac{\|y_i - \tilde{x}_k\|}{\sigma}\right)$

где $\sigma$ — это дистанция влияния точек ядра, которая выбирается в соответствии с плотностью входных данных.

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

$(F \ast g)(x) = \sum_{x_i \in N_x} g(x_i - x) f_i$

Это может выглядеть немного запутанно и непонятно, но на примере будут видны все этапы.

Простой пример расчета в коде

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

import numpy as np

# Определяем координаты центральной точки
x = np.array([0.0, 0.0, 0.0])

# Координаты трех соседних точек
neighbors = np.array([
    [0.5, 0.0, 0.0],
    [0.2, 0.1, 0.0],
    [-0.3, 0.0, 0.1]
])

# Координаты четырех точек ядра
kernel_points = np.array([
    [0.1, 0.1, 0.1],
    [-0.1, -0.1, 0.1],
    [0.1, -0.1, -0.1],
    [-0.1, 0.1, -0.1]
])

# Матрицы весов для каждой точки ядра (пример с двумерным пространством признаков для наглядности)
weight_matrices = np.array([
    [[2.0, 0.5], [1.5, 1.0]],
    [[1.0, 0.5], [2.0, 1.5]],
    [[1.5, 1.0], [1.0, 2.0]],
    [[0.5, 1.5], [2.0, 0.5]]
])

# Сигма для функции корреляции
sigma = 0.5

# Функция для вычисления корреляции между точкой ядра и соседом
def h(yi, xk, sigma):
    distance = np.linalg.norm(yi - xk)
    return max(0, 1 - distance / sigma)

# Вычисляем корреляцию для каждого соседа относительно каждой точки ядра
correlations = np.zeros((neighbors.shape[0], kernel_points.shape[0]))

for i, neighbor in enumerate(neighbors):
    yi = neighbor - x
    for k, kernel_point in enumerate(kernel_points):
        correlations[i, k] = h(yi, kernel_point, sigma)

# Признаки для каждой из соседних точек
features = np.array([
    [1, 2],
    [3, 4],
    [5, 6]
])

# Вычисляем вклад каждой точки ядра в признаки центральной точки
new_features = np.zeros(2)
for i, neighbor in enumerate(neighbors):
    neighbor_contribution = np.zeros(2)
    for k in range(kernel_points.shape[0]):
        neighbor_contribution += correlations[i, k] * weight_matrices[k].dot(features[i])
    new_features += neighbor_contribution

Deformable kernel

Мы рассмотрели так называемые жесткие ядра (rigid kernel). Они хороши, но хочется, что бы наша сеть могла динамично подстраиваться под каждую уникальную структуру данных, с которой она сталкивается. Это и есть суть деформируемых ядер: они не статичны. Далее опишу математический аппарат и логику уже без примера расчета в коде.

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

$(F \ast g)(x) = \sum_{x_i \in N_x} g_{\text{deform}}(x - x_i, \Delta(x)) f_i$

Где $g_{deform}$ — это деформируемая функция ядра, в которой смещение модифицирует позиции ядер, обеспечивая их способность к адаптации:

$g_{\text{deform}}(y_i, \Delta(x)) = \sum_{k<K} h(y_i, \tilde{x}_k + \Delta(x)) W_k$

Регуляризация этих смещений производится через сумму притягивающего и отталкивающего членов (перевод корректный по смыслу, но не много режет слух. Возможно есть более подходящие варианты). Притягивающий член $L_{fit}$ помогает каждой точке ядра поддерживать нужное расстояние до соседей, а отталкивающий член $L_{rep}$ предотвращает слишком тесное сближение точек ядра, что могло бы привести к перекрытию их областей влияния:

$L_{\text{reg}} = \sum_x L_{\text{fit}}(x) + L_{\text{rep}}(x)$ $L_{\text{fit}}(x) = \sum_{y_i} \min_{k<K} \left( \frac{\|y_i - (\tilde{x}_k + \Delta(x))\|}{\sigma} \right)^2$ $L_{\text{rep}}(x) = \sum_{k<K} \sum_{l \neq k} h(\tilde{x}_k + \Delta(x), \tilde{x}_l + \Delta(x))^2$

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

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

Наглядный результат можно видеть ниже:

Положение элементов ядра в Rigid Kernel и Deformable Kernel

IoU при использовании Rigid Kernel и Deformable Kernel

Пример сегментации с KPConv

Выводы

Надеюсь, я смог понятно изложить принципы KPConv. В процессе моего исследования я составил докер-контейнер с KPConv, обучил модель на данных KITTI и затем на собственном наборе данных, что дало отличные результаты. Конечно, всё началось с задачи совмещения облаков точек, и вот что интересно: в то время как я изучал KPConv, мой коллега нашёл решение с помощью алгоритма NDT (Normal distributions transform), который удивил своей простотой. Скорее всего, моя следующая статья будет посвящена именно ему, и я планирую реализовать его с нуля, что бы понять максимально полно данный алгоритм.

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

P.S. Следующая задача будет сегментирование облака точек, так что не зря изучил данный слой и собрал всё в один докер-контейнер :-)

KPConv на русском: Свертки для point cloud