Pull to refresh

Быстрое вычисление точной 3D карты расстояний с использованием технологии CUDA

Algorithms *
Карта расстояний (Distance Map) — это объект, позволяющий быстро получить расстояние от заданной точки до определенной поверхности. Обычно представляет собой матрицу значений расстояний для узлов с фиксированным шагом. Часто используется в играх для определения «попадания» в игрока или предмет, и для оптимизационных задач по совмещению объектов: расположить объекты максимально близко друг к другу, но так, чтобы они не пересекались. В первом случае качество карты расстояний (то есть точность значений в узлах) не играет большой роли. Во втором — от нее могут зависеть жизни (в ряде приложений, связанных с нейрохирургией). В этой статье я расскажу как можно достаточно точно обсчитать карту расстояний за разумное время.

Основные объекты

Пусть у нас имеется некоторая поверхность S, заданная набором вокселей. Координаты вокселей будем считать на регулярной сетке (т.е. шаги по X, Y и Z одинаковы).

Требуется вычислить карту расстояний M[X,Y,Z] для всех вокселей, которые лежат в кубе, содержащем поверхность; M[x,y,z] = d((x,y,z), S) = min( d((x,y,z), (Snx, Sny, Snz) ) = sqrt( min( (x-Snx)2 + (y-Sny)2 + (z-Snz)2) ). Последнее равенство верно только для регулярной сетки, остальное должно быть очевидно.

Откуда берутся данные


Для нашей задачи данные поступают с аппарата МРТ. Напомню, что мы работаем с 3D изображениями, так что весь 3D образ представляет собой набор плоских картинок для разных срезов по Z, которые я, разумеется не могу все здесь представить; их немногим меньше, чем разрешение по X и Y. Типичный размер томографии порядка 625x625x592 вокселей.


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

Прямой перебор

Прикинем для начала, почему бы не заполнить карту расстояний значениями «по определению» — минимумами расстояний до точек границы. Всего точек карты: X*Y*Z, точек поверхности порядка max(X,Y,Z)2. Вспомним размеры исходных данных и получим что-то порядка 592*6254 арифметико-логических операций. Калькулятор подскажет, что это более чем 90000 миллиардов. Мягко говоря, многовато, отложим пока прямой перебор.

Динамическое программирование

Напрашивается использование того, что данные представляются в трехмерном массиве; можно как-то рассчитать значение в каждой точке, используя ее соседей, и разумеется мы не первые, кто этим занимался. Несколько отталкивает факт, что держать в памяти структуру размером 592*625*625*sizeof(float), что около 1 Гигабайта, несколько ресурсоемко. Ну и ладно, забудем пока что томография может сниматься и с двойным разрешением (еще в 8 раз больше) — 640Кбайт давно уже никому не хватает.

Самый тривиальный метод заполнения (city-block), похожий на алгоритм для обхода лабиринта, отрабатывает за считанные секунды, но максимальная ошибка на трехмерной сетке может составить 42 шага сетки для максимально обсчитываемого расстояния в 50 шагов. Насколько это плохо? Полный провал.

Единственный алгоритм из класса (central point) дает почти удовлетворительную точность (ошибка не более 0.7 шага сетки), но работает минуты.

Мы пытались найти компромисс, реализовав chessboard distance (описание в той же статье, что и central point), но итоговая точность не могла устроить, а время его работы — десятки секунд.

Картинка — ссылка на полноразмерную, но даже на превью видно «странное» поведение достаточно далеко от точек поверхности, и особенно если нужно обсчитывать еще и внутренние объекты. Слева chessboard, а справа то, что было получено прямым перебором. И еще раз напомню, для тех, кто смотрит только картинки — речь идет о 3D изображениях, поэтому странные «проколы» внутри объяснимы — они вызваны близостью к поверхности на следующих или предыдущих слайдах по Z. В общем, провал.

Опять прямой перебор

Да, слово CUDA на картинке как бы намекает. Может быть 90000 миллиардов — не так уж и много, а может это число можно как-то уменьшить, без потери точности. В любом случае, middle-end видеокарта справляется с примитивными вычислительными операциями в 80-120 раз быстрее процессора аналогичного класса. Попробуем.

Да и у прямого перебора есть существенное преимущество — нет необходимости держать всю карту расстояний в памяти, и как следствие — линейная масштабируемость.

Сокращение перебора

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

Критерий корректности — абсолютная точность вблизи границы (на расстоянии до 10 (GUARANTEED_DIST) шагов сетки), хорошая точность на расстоянии до 36 (TARGET_MAX_DIST) шагов, и визуальная адекватность (без артефактов — зачем смущать хирургов). Числа 10 и 36 взялись непосредственно из нашей задачи (у нас вообще это значения в миллиметрах и сетка миллиметровая).

Кубический индекс

Для каждой точки хотелось бы определить — насколько точно ее значение нужно вычислять. Очевидно, что размер такой структуры совпадет с размером исходного изображения, а это много, и считаться будет долго.

Будем объединять точки в «кубики». Пусть размер кубика — 32 (CUBE_SIZE), тогда индекс будет иметь размерности 20 x 20 x 19.

Маркируем кубик по типу «лучшей» точки, лежащей в нем. Если хоть одна точка требует точного обсчета — считаем весь кубик точно. В худшем случае — не считаем вообще.

Как быстро промаркировать кубики? Начнем с точек границы: прибавим к каждой из них смещение GUARANTEED_DIST во всех 27 направлениях ((-1,0,1)x(-1,0,1)x(-1,0,1)), получим индекс соответствующего кубика, делением на 32 (CUBE_SIZE), и пометим его как «точный» — ведь в нем лежит по крайней мере одна точка, требующая точного обсчета. Задача: объяснить почему этого достаточно, и все кубики содержащие точки для точного обсчета будут помечены верно.

На всякий случай довольно тривиальный код, который это выполняет:

  for (size_t t = 0; t < borderLength; t++)
  {
    for (int dz = -1; dz <= 1; dz++)
      for (int dy = -1; dy <= 1; dy++)
        for (int dx = -1; dx <= 1; dx++)
        {
          float x = border[t].X + (float)dx * GUARANTEED_DIST;
          float y = border[t].Y + (float)dy * GUARANTEED_DIST;
          float z = border[t].Z + (float)dz * GUARANTEED_DIST;
          if (x >= 0 && y >= 0 && z >= 0)
          {
            size_t px = round(x / CUBE_SIZE);
            size_t py = round(y / CUBE_SIZE);
            size_t pz = round(z / CUBE_SIZE);
            size_t idx = px + dim_X * (py + skip_Y * pz);
            if (idx < dim_X * dim_Y * dim_Z)
            {
              markedIdx[idx] = true;
            }
          }
        }
  }



На практических данных эта идея позволила сократить порядок перебора «всего» в 4-5 раз. Все равно получается очень много, но идея с кубиками нам еще пригодится.

Кубический контейнер

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

Раскидаем все воксели границы S по кубикам (про индексные кубики с настоящего момента забудем; мы их построили, поняли как применять, и теперь больше с ними не работаем). В каждый кубик будем класть те точки границы, которые достижимы (лежат на расстоянии TARGET_MAX_DIST и меньше). Для этого достаточно посчитать центр кубика, и отложить от него расстояние sqrt(3)/2*CUBE_SIZE (это диагональ кубика) + TARGET_MAX_DIST. Если точка достижима из вершины кубика или любой точки стороны, то она достижима из его внутренности.

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

Оптимизации:
  • Первая (вынужденная) — размер кубика подбирается так, чтобы итоговое число «размноженных» точек не было больше максимального объема памяти на них (мы брали 512 мегабайт). Подобрать размер кубика можно быстро (немного математики).
  • Вторая (разумная) — использовать индексы на начальное и конечное размещение точек для каждого кубика, эти индексы могут пересекаться, и за счет грамотного их вычисления можно сэкономить еще в 2 раза по объему памяти.

Я сознательно не буду приводить точные алгоритмы для обоих пунктов (а вдруг вы реализуете их лучше и станете нашими конкурентами? ;) ), там слов больше, чем идеи.

В целом кубики сокращают порядок перебора в X/CUBE_SIZE * Y/CUBE_SIZE * Z/CUBE_SIZE раз, но надо помнить, что уменьшение размера кубика потребует существенно больше памяти, и пороговое значение для наших изображений оказывалось в районе 24-32.

Таким образом, задача свелась к порядка 100-200 миллиардов вещественных операций. Теоретически, это вычислимо на видеокарте за десятки секунд.

Ядро CUDA

Оно получилось такое простое и красивое, что не поленюсь его показать:

__global__ void minimumDistanceCUDA( /*out*/ float *result,
                   /*in*/ float *work_points,
                   /*in*/ int  *work_indexes,
                   /*in*/ float *new_border)
{        
  __shared__ float sMins[MAX_THREADS]; // max threads count
  
  for (int i = blockIdx.x; i < BLOCK_SIZE_DIST; i += gridDim.x)
  {
    sMins[threadIdx.x] = TARGET_MAX_DIST*TARGET_MAX_DIST;

    int startIdx = work_indexes[2*i];
    int endIdx  = work_indexes[2*i+1];    
    
    float x = work_points[3*i];
    float y = work_points[3*i+1];
    float z = work_points[3*i+2];
    
    // main computational entry
    for (int j = startIdx + threadIdx.x; j < endIdx; j += blockDim.x)
    {        
      float dist = (x - new_border[3*j] )*(x - new_border[3*j] )
            + (y - new_border[3*j+1])*(y - new_border[3*j+1])
            + (z - new_border[3*j+2])*(z - new_border[3*j+2]);
      if (dist < sMins[threadIdx.x])
        sMins[threadIdx.x] = dist;
    }

    __syncthreads();

    if (threadIdx.x == 0)
    {
      float min = sMins[0];      
      for (int j = 1; j < blockDim.x; j++)
      {
        if (sMins[j] < min)
          min = sMins[j];
      }
      result[i] = sqrt(min);
    }        

    __syncthreads();
  }
}


Мы загружаем блок координат точек distance map для перебора (work_points), соответствующие начальные и конечные индексы для перебора точек границы (work_indexes) и границу, оптимизированную алгоритмом с «кубиками» (new_border). Результаты забираем в одинарной вещественной точности.

Собственно, сам код ядра, что приятно очень понятен, даже без тонкостей CUDA. Главной его особенностью является использование переменных gridDim.x и threadIdx.x обозначающих индексы текущего потока из набора запущенных потоков на видеокарте — все делают одно и то же, но с разными точками, а потом результаты синхронизируются и аккуратно записываются.

Осталось только грамотно организовать вычисление work_points на CPU (но это же так просто!) и запустить вычислитель:
#define GRID       512   // calculation grid size
#define THREADS     96   // <= MAX_THREADS, optimal for this task
minimumDistanceCUDA<<<GRID, THREADS>>>(dist_result_cuda, work_points_cuda, work_indexes_cuda, new_border_cuda);


На Core Duo E8400 и видеокарте GTX460 средняя томография обрабатывается за время в пределах минуты, расход памяти лимитирован 512Мб. Распределение процессорного времени составляет примерно процентов 20 на CPU и файловые операции, и оставшиеся 80 на вычисление минимумов на видеокарте.

А вот так выглядит картинка посчитанная с «загрублением» далеких точек:



Как-то так, надеюсь что-нибудь Вам пригодится. Большая просьба, не копировать картинки и не делать ре-публикации без моего разрешения (автор метода я, хотя «методом» это сложно назвать).

Удачного дня!
Tags:
Hubs:
Total votes 51: ↑50 and ↓1 +49
Views 2.6K
Comments Comments 21