Измерение расстояния до объектов с помощью RealSense D435

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

Результаты обнаружения объекта могут быть дополнены информацией о том, насколько далеко расположен данный объект. Задачу измерения расстояния можно решить с помощью камеры глубины Intel RealSense D435, измеряющей глубину в каждой точке.

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

Требования

Камера RealSense

Для измерения расстояния нам понадобится камера глубины RealSense D435. Для данного примера устанавливать RealSense SDK нет необходимости.

Библиотека pyrealsense2

Весь код, для простоты демонстрации, мы будем писать на Python 3.7, так что нам понадобится пакет pyrealsense2.

pip install pyrealsense2

Библиотека OpenCV

Также, нам понадобится библиотека OpenCV (подойдет любая версия, начиная с 3.4). Мануал по установке библиотеки можно посмотреть на официальном сайте OpenCV.

Вместе с OpenCV установится еще одна необходимая библиотека – numpy.

Обнаружение объектов

Первым делом необходимо выделить объекты на видео. Для обнаружения объектов в режиме реального времени отлично подойдут так называемые one-stage модели (например, Retina Net, SSD, YOLO), которые выигрывают в скорости работы.

Для простоты эксперимента возьмем обученную модель SSDLite Mobilenet v2. Вместо нее может быть использована любая другая модель, которая возвращает координаты объекта в том или ином виде.
Скачать архив с файлами модели можно в официальном репозитории TensorFlow.

Архив содержит замороженный граф вычислений в бинарном формате .pb, а также различные конфиги. Но для использования модели в OpenCV необходимо также иметь граф вычислений в текстовом формате .pbtxt. В архиве он отсутствует, поэтому его нужно сгенерить вручную.

Копируем два скрипта из репозитория OpenCV:

1. github.com/opencv/opencv/blob/master/samples/dnn/tf_text_graph_ssd.py
2. github.com/opencv/opencv/blob/master/samples/dnn/tf_text_graph_common.py

Затем, выполняем следующую команду, указав нужные пути:

python tf_text_graph_ssd.py --input /path/to/model.pb --config /path/to/example.config --output /path/to/graph.pbtxt

На выходе получаем файл graph.pbtxt.

Имеющиеся файлы модели кладем в папку с проектом. В моем случае, это папка models/.

Разработка

Теперь все готово к разработке. Начнем писать код.

Импортируем необходимые библиотеки и делаем заглушку на метод draw_predictions(), который понадобится нам позже:

import pyrealsense2 as rs
import numpy as np
import cv2 as cv

def draw_predictions(pred_img, color_img, depth_image):
    pass

Создаем основный объекты:

pipeline = rs.pipeline() # <- Объект pipeline содержит методы для взаимодействия с потоком
config = rs.config() # <- Дополнительный объект для хранения настроек потока
colorizer = rs.colorizer() # <- Пригодится для отрисовки цветной карты глубины

spatial = rs.spatial_filter()
spatial.set_option(rs.option.holes_fill, 3)

Инициализируем ssd-модель. Тут нам пригодится .pbtxt файл, который мы сгененировали выше:

# Инициализируем модель
model = cv.dnn.readNetFromTensorflow(
    "models/ssdlite_mobilenet_v2_coco_2018_05_09/frozen_inference_graph.pb", 
    "models/ssdlite_mobilenet_v2_coco_2018_05_09/graph.pbtxt"
)

Запускаем потоки камеры:

# Запускаем 2 потока: RGB и depth
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
profile = pipeline.start(config)

обратите внимание, что мы захватываем цветной поток в формате BGR, так как именно этот формат использует OpenCV по-умолчанию.

Заводим цикл, в которым будем последовательно захватывать и обрабатывать фреймы:

cv.namedWindow("RealSense object detection", cv.WINDOW_AUTOSIZE)
try:
    while True:
        # Ждем захват фреймов для "цвета" и "глубины"
        frames = pipeline.wait_for_frames()
        depth_frame = frames.get_depth_frame()
        color_frame = frames.get_color_frame()
        if not depth_frame or not color_frame:
            continue

Для дальнейшей обработки фреймов представляем их в виде numpy-массивов:

        # Конвертируем фреймы в numpy-массивы
        depth_image = np.asanyarray(depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())

Применяем ssd-модель для обнаружения:

        # Обнаружение объектов
        model.setInput(cv.dnn.blobFromImage(color_image, size=(300, 300), swapRB=True, crop=False))
        pred = model.forward()
        draw_predictions(pred, color_image, depth_image)

Рисуем результат в окне:

	# Переводим изображение глубины в цвет
        colorized_depth = cv.applyColorMap(cv.convertScaleAbs(depth_image, alpha=0.03), cv.COLORMAP_JET)
        # Соединяем и показываем изображения
        images = np.hstack((color_image, colorized_depth))
        cv.imshow("RealSense object detection", images)

Пишем условие выхода из цикла:

	# Выйти при нажатии 'ESC' или 'q'
        key = cv.waitKey(1) & 0xFF
        if (key == 27) or (key == ord('q')):
            cv.destroyWindow("RealSense object detection")
            break        

Закрываем поток:

finally:
    pipeline.stop()

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

def draw_predictions(pred_img, color_img, depth_image):
    for detection in pred_img[0,0,:,:]:
        score = float(detection[2])
        # Рисуем рамку только при уверенности модели в обнаружении выше чем на 50%
        if score > 0.5:
            left = detection[3] * color_img.shape[1]
            top = detection[4] * color_img.shape[0]
            right = detection[5] * color_img.shape[1]
            bottom = detection[6] * color_img.shape[0]
            cv.rectangle(color_img, (int(left), int(top)), (int(right), int(bottom)), (210, 230, 23), 2)
            
            # Измеряем расстояние до объекта
            depth = depth_image[int(left):int(right), int(top):int(bottom)].astype(float)
            depth_scale = profile.get_device().first_depth_sensor().get_depth_scale()
            depth = depth * depth_scale
            dist,_,_,_ = cv.mean(depth)
            dist = round(dist, 1)
            cv.putText(color_img, "dist: "+str(dist)+"m", (int(left), int(top)-5), cv.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)

Полный код

import pyrealsense2 as rs
import numpy as np
import cv2 as cv

def draw_predictions(pred_img, color_img, depth_image): # <- Метод для отрисовки рамки    
    for detection in pred_img[0,0,:,:]:
        score = float(detection[2])
        # Рисуем рамку только при уверенности модели в обнаружении выше чем на 50%
        if score > 0.5:
            left = detection[3] * color_img.shape[1]
            top = detection[4] * color_img.shape[0]
            right = detection[5] * color_img.shape[1]
            bottom = detection[6] * color_img.shape[0]
            cv.rectangle(color_img, (int(left), int(top)), (int(right), int(bottom)), (210, 230, 23), 2)
            
            # Измеряем расстояние до объекта
            depth = depth_image[int(left):int(right), int(top):int(bottom)].astype(float)
            depth_scale = profile.get_device().first_depth_sensor().get_depth_scale()
            depth = depth * depth_scale
            dist,_,_,_ = cv.mean(depth)
            dist = round(dist, 1)
            cv.putText(color_img, "dist: "+str(dist)+"m", (int(left), int(top)-5), cv.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1)


pipeline = rs.pipeline() # <- Объект pipeline содержит методы для взаимодействия с потоком
config = rs.config() # <- Дополнительный объект для хранения настроек потока
colorizer = rs.colorizer() # <- Пригодится для отрисовки цветной карты глубины

# Инициализируем модель
model = cv.dnn.readNetFromTensorflow(
    "models/ssdlite_mobilenet_v2_coco_2018_05_09/frozen_inference_graph.pb", 
    "models/ssdlite_mobilenet_v2_coco_2018_05_09/graph.pbtxt"
)

# Запускаем 2 потока: RGB и depth
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
profile = pipeline.start(config)

cv.namedWindow("RealSense object detection", cv.WINDOW_AUTOSIZE)
try:
    while True:
        # Ждем захват фреймов для "цвета" и "глубины"
        frames = pipeline.wait_for_frames()
        depth_frame = frames.get_depth_frame()
        color_frame = frames.get_color_frame()
        if not depth_frame or not color_frame:
            continue

        # Конвертируем фреймы в numpy-массивы
        depth_image = np.asanyarray(depth_frame.get_data())
        color_image = np.asanyarray(color_frame.get_data())
        
        # Обнаружение объектов
        model.setInput(cv.dnn.blobFromImage(color_image, size=(300, 300), swapRB=True, crop=False))
        pred = model.forward()
        draw_predictions(pred, color_image, depth_image)

        # Переводим изображение глубины в цвет
        colorized_depth = cv.applyColorMap(cv.convertScaleAbs(depth_image, alpha=0.03), cv.COLORMAP_JET)
        # Соединяем и показываем изображения
        images = np.hstack((color_image, colorized_depth))
        cv.imshow("RealSense object detection", images)
        
        # Выйти при нажатии 'ESC' или 'q'
        key = cv.waitKey(1) & 0xFF
        if (key == 27) or (key == ord('q')):
            cv.destroyWindow("RealSense object detection")
            break

finally:
    pipeline.stop()

Заключение

Таким образом, у нас имеется простенький скрипт для оценки расстояния до выделенного объекта. Алгоритм может ошибаться примерно на 50 см, но в целом работает хорошо в пределах четырех метров. Для увеличения точности замера можно применить к карте глубин дополнительные фильтры, заложенные в pyrealsense2, увеличивающие качество изображения, либо модифицировать сам алгоритм вычисления глубины (напр. вычислять взвешенное среднее или замерять расстояние в одной центральной области). Как видите задача не сложная, но интересная.