Семантическая сегментация с Open3D-ML, PyTorch Backend и пользовательским набором данных

Примечание. Инструкции по загрузке, запуску и устранению неполадок кода, представленного в этой статье, приведены в конце.

В рамках своих экспериментов с Open3D-ML для облаков точек я написал статьи, объясняющие, как установить эту библиотеку с поддержкой Tensorflow и PyTorch. Чтобы протестировать установку, я объяснил, как запустить простой скрипт Python для визуализации помеченного набора данных для семантической сегментации под названием SemanticKITTI. В этой статье я расскажу о шагах, которые я выполнил, чтобы сделать вывод о любом облаке точек, включая тестовую часть SemanticKITTI, а также о моем частном наборе данных.

В остальной части этой статьи предполагается, что вы успешно установили и протестировали Open3D-ML с серверной частью PyTorch, следуя моей предыдущей статье. Это также означает, что вы загрузили набор данных SemanticKITTI. Чтобы запустить модель семантической сегментации для немаркированных данных, вам необходимо загрузить конвейер Open3D-ML. Конвейер будет состоять из модели семантической сегментации, набора данных и, возможно, других этапов предварительной/постобработки. Open3D-ML поставляется с модулями и файлами конфигурации для простой загрузки и запуска популярных конвейеров.

Чтобы сделать вывод о новых облаках точек, мы будем использовать популярную модель под названием RandLA-Net, представленную в статье 2019 года под названием RandLA-Net: эффективная семантическая сегментация крупномасштабных облаков точек. Удобно, что Open3D-ML имеет реализацию этого метода и имеет конфигурации для загрузки и запуска такого метода в наборе данных SemanticKITTI без особых усилий.

Чтобы загрузить файл конфигурации, нам нужен следующий код, заменив /path/to/Open3D/ на путь, по которому вы клонировали репозиторий Open3D при установке.

# Load an ML configuration file
cfg_file = "/path/to/Open3D/build/Open3D-ML/ml3d/configs/randlanet_semantickitti.yml"
cfg = _ml3d.utils.Config.load_from_file(cfg_file)

Далее мы создадим модель RandLANet, используя объект конфигурации, и добавим пути к набору данных SemanticKITTI, а также к нашему пользовательскому набору данных. Обязательно замените /path/to/save/dataset/SemanticKitti/ на путь, по которому вы сохранили данные SemanticKITTI при установке Open3D-ML. На данный момент пользовательский набор данных указывает на некоторые из моих личных облаков точек, собранных с помощью моего робота и представленных в репозитории, сопровождающем эту статью.

# Load the RandLANet model
model = ml3d.models.RandLANet(**cfg.model)
# Add path to the SemanticKitti dataset and your own custom dataset
cfg.dataset['dataset_path'] = '/path/to/save/dataset/SemanticKitti/'
cfg.dataset['custom_dataset_path'] = './pcds'

Следующим шагом является загрузка наборов данных. Для загрузки набора данных SementicKITTI в Open3D-ML есть удобные вспомогательные классы и методы.

# Load the datasets
dataset = ml3d.datasets.SemanticKITTI(cfg.dataset.pop('dataset_path', None), **cfg.dataset)
custom_dataset = load_custom_dataset(cfg.dataset.pop('custom_dataset_path', None))

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

def load_custom_dataset(dataset_path):
 print("Loading custom dataset")
 pcd_paths = glob.glob(dataset_path+"/*.pcd")
 pcds = []
 for pcd_path in pcd_paths:
  pcds.append(o3d.io.read_point_cloud(pcd_path))
 return pcds

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

# Create the ML pipeline
pipeline = ml3d.pipelines.SemanticSegmentation(model, dataset=dataset, device="gpu", **cfg.pipeline)
# Download the weights.
ckpt_folder = "./logs/"
os.makedirs(ckpt_folder, exist_ok=True)
ckpt_path = ckpt_folder + "randlanet_semantickitti_202201071330utc.pth"
randlanet_url = "https://storage.googleapis.com/open3d-releases/model-zoo/randlanet_semantickitti_202201071330utc.pth"
if not os.path.exists(ckpt_path):
    cmd = "wget {} -O {}".format(randlanet_url, ckpt_path)
    os.system(cmd)
# Load the parameters of the model.
pipeline.load_ckpt(ckpt_path=ckpt_path)

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

# Get one test point cloud from the SemanticKitti dataset
pc_idx = 58 # change the index to get a different point cloud
test_split = dataset.get_split("test")
data = test_split.get_data(pc_idx)
# run inference on a single example.
# returns dict with 'predict_labels' and 'predict_scores'.
result = pipeline.run_inference(data)

Экземпляр данных Point Cloud в наборе данных SemanticKITTI загружается как словарь Python, содержащий ключи «point», «feat» и «label». Последние два имеют None и массив Numpy, заполненный нулями в качестве значений соответственно, и не используются во время логического вывода. Ключ point связан с массивом Numpy, содержащим координаты x, y и z точек LiDAR. Чтобы визуализировать результат вывода с помощью визуализатора Open3D, нам нужно создать объект «Облако точек» из «точки» части словаря, а затем раскрасить точки с помощью меток, возвращенных выводом.

# Create a pcd to be visualized 
pcd = o3d.geometry.PointCloud()
xyz = data["point"] # Get the points
pcd.points = o3d.utility.Vector3dVector(xyz)
# Get the color associated with each predicted label
colors = [COLOR_MAP[clr] for clr in list(result['predict_labels'])] 
pcd.colors = o3d.utility.Vector3dVector(colors) # Add color data to the point cloud
# Create visualization
custom_draw_geometry(pcd)

# Class colors, RGB values as ints for easy reading
COLOR_MAP = {
    0: (0, 0, 0),
    1: (245, 150, 100),
    2: (245, 230, 100),
    3: (150, 60, 30),
    4: (180, 30, 80),
    5: (255, 0., 0),
    6: (30, 30, 255),
    7: (200, 40, 255),
    8: (90, 30, 150),
    9: (255, 0, 255),
    10: (255, 150, 255),
    11: (75, 0, 75),
    12: (75, 0., 175),
    13: (0, 200, 255),
    14: (50, 120, 255),
    15: (0, 175, 0),
    16: (0, 60, 135),
    17: (80, 240, 150),
    18: (150, 240, 255),
    19: (0, 0, 255),
}
# Convert class colors to doubles from 0 to 1, as expected by the visualizer
for label in COLOR_MAP:
 COLOR_MAP[label] = tuple(val/255 for val in COLOR_MAP[label])

def custom_draw_geometry(pcd):
 vis = o3d.visualization.Visualizer()
 vis.create_window()
 vis.get_render_option().point_size = 2.0
 vis.get_render_option().background_color = np.asarray([1.0, 1.0, 1.0])
 vis.add_geometry(pcd)
 vis.run()
 vis.destroy_window()

# Get one test point cloud from the custom dataset
pc_idx = 2 # change the index to get a different point cloud
data, pcd = prepare_point_cloud_for_inference(custom_dataset[pc_idx])
# Run inference
result = pipeline.run_inference(data)
# Colorize the point cloud with predicted labels
colors = [COLOR_MAP[clr] for clr in list(result['predict_labels'])]
pcd.colors = o3d.utility.Vector3dVector(colors)
# Create visualization
custom_draw_geometry(pcd)

def prepare_point_cloud_for_inference(pcd):
 # Remove NaNs and infinity values
 pcd.remove_non_finite_points()
 # Extract the xyz points
 xyz = np.asarray(pcd.points)
 # Set the points to the correct format for inference
 data = {"point":xyz, 'feat': None, 'label':np.zeros((len(xyz),), dtype=np.int32)}
return data, pcd

git clone https://github.com/carlos-argueta/open3d_experiments.git

cd open3d_experiments
python3 semantic_torch.py

pip install tensorboard