[open3d] 9. 3d semantic segmentation model 추론

[Sungryul Lee]

1. 3d segmentation model을 이용한 추론

open3d-ml 라이브러리를 이용하여 semantic segmentation 모델의 추론을 수행하는 예제

사용할 semantic segmentation 모델 : kpconv, randlanet -> 2가지 모델만 지원함

위 2가지 모델의 훈련된 가중치를 이용하여 모델의 추론만 수행할 것임

테스트 데이터는 SemanticKITTI 데이터셋의 test용 데이터를 이용할 것임

SemanticKITTI  -> https://semantic-kitti.org/dataset.html#download

실행환경 ; wsl2-ubuntu24.04, anaconda 가상환경, open3d-ml 예제활용

데이터셋은 윈도우즈에서 다운받아 저장하고 wsl2-ubuntu24.04에서 공유가능, 예를 들어 윈도우즈 c드라이브, 리눅스 /mnt/c으로 공유됨 -> 따라서 윈도우즈에서 다운받은 데이터셋을 리눅스 디렉토리로 복사할 필요없음

dataset_path ->  /mnt/d/Users/2sungryul/Dropbox/Work/Dataset/SemanticKITTI/data_odometry_velodyne

그 아래에 다음 형태로 서브디렉토리가 저장되어 있어야함 -> dataset_path\dataset\sequences\xx\velodyne, labels

-> 라이브러리 코드에서 위와 같은 형태로 서브 디렉토리를 가정하고 데이터를 접근함

클래스 레이블 정보

kitti_labels = {

    0: 'unlabeled',

    1: 'car',

    2: 'bicycle',

    3: 'motorcycle',

    4: 'truck',

    5: 'other-vehicle',

    6: 'person',

    7: 'bicyclist',

    8: 'motorcyclist',

    9: 'road',

    10: 'parking',

    11: 'sidewalk',

    12: 'other-ground',

    13: 'building',

    14: 'fence',

    15: 'vegetation',

    16: 'trunk',

    17: 'terrain',

    18: 'pole',

    19: 'traffic-sign'

}

2. 소스코드

# Running a pretrained model for semantic segmentation

import os

import sys

from os.path import exists, join, dirname

import open3d.ml as _ml3d

import open3d.ml.torch as ml3d

import numpy as np

"""

    test_split: ['11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21']

    training_split: ['00', '01', '02', '03', '04', '05', '06', '07', '09', '10']

    all_split: ['00', '01', '02', '03', '04', '05', '06', '07', '09',

                '08', '10', '11', '12', '13', '14', '15', '16', '17', '18', '19', '20', '21']

    validation_split: ['08']

"""

example_dir = os.path.dirname(os.path.realpath(__file__))

def get_torch_ckpts():

    kpconv_url = "https://storage.googleapis.com/open3d-releases/model-zoo/kpconv_semantickitti_202009090354utc.pth"

    randlanet_url = "https://storage.googleapis.com/open3d-releases/model-zoo/randlanet_semantickitti_202201071330utc.pth"

    ckpt_path_r = example_dir + "/vis_weights_{}.pth".format('RandLANet')

    if not exists(ckpt_path_r):

        cmd = "wget {} -O {}".format(randlanet_url, ckpt_path_r)

        os.system(cmd)

    ckpt_path_k = example_dir + "/vis_weights_{}.pth".format('KPFCNN')

    if not exists(ckpt_path_k):

        cmd = "wget {} -O {}".format(kpconv_url, ckpt_path_k)

        print(cmd)

        os.system(cmd)

    return ckpt_path_r, ckpt_path_k

vis_points = []

def pred_custom_data(index, pcs, pipeline_r, pipeline_k):

    #vis_points = []

    name = "{}".format(index)

    results_r = pipeline_r.run_inference(pcs)

    pred_label_r = (results_r['predict_labels'] + 1).astype(np.int32)

    # Fill "unlabeled" value because predictions have no 0 values.

    pred_label_r[0] = 0

    results_k = pipeline_k.run_inference(pcs)

    pred_label_k = (results_k['predict_labels'] + 1).astype(np.int32)

    # Fill "unlabeled" value because predictions have no 0 values.

    pred_label_k[0] = 0

    label = pcs['label']

    pts = pcs['point']

    vis_d = {

        "name": name,

        "points": pts,

        "labels": label,

        "pred": pred_label_k,

    }

    vis_points.append(vis_d)

    vis_d = {

        "name": name + "_randlanet",

        "points": pts,

        "labels": pred_label_r,

    }

    vis_points.append(vis_d)

    vis_d = {

        "name": name + "_kpconv",

        "points": pts,

        "labels": pred_label_k,

    }

    vis_points.append(vis_d)

    #return vis_points

kitti_labels = ml3d.datasets.SemanticKITTI.get_label_to_names()

v = ml3d.vis.Visualizer()

lut = ml3d.vis.LabelLUT()

for val in sorted(kitti_labels.keys()):

    lut.add_label(kitti_labels[val], val)

v.set_lut("labels", lut)

v.set_lut("pred", lut)

# load pretrained weights -> 현재 작업디렉토리에 가중치파일이(pth)이 존재하면 그것을 사용하고 아니면 인터넷저장소에서 다운로드 해준다

ckpt_path_r, ckpt_path_k = get_torch_ckpts()

# RandLaNet 모델, 파이프라인 객체생성후 가중치 파일로 초기화

model = ml3d.models.RandLANet(ckpt_path=ckpt_path_r)

pipeline_r = ml3d.pipelines.SemanticSegmentation(model,device="gpu")

pipeline_r.load_ckpt(model.cfg.ckpt_path)

# KPFCNN 모델, 파이프라인 객체생성후 가중치 파일로 초기화

model = ml3d.models.KPFCNN(ckpt_path=ckpt_path_k)

pipeline_k = ml3d.pipelines.SemanticSegmentation(model,device="gpu")

pipeline_k.load_ckpt(model.cfg.ckpt_path)

# load dataset

dataset_path = "/mnt/d/Users/2sungryul/Dropbox/Work/Dataset/SemanticKITTI/data_odometry_velodyne"

dataset = ml3d.datasets.SemanticKITTI(dataset_path)

# load train dataset -> sequences 00~07, 09, 10

train_split = dataset.get_split("train")

len = train_split.__len__()

print('train:',len)

train_data = train_split.get_data(0)

print(train_data.keys())

#train_data['feat']=np.squeeze(train_data['feat'])

train_data['feat']=None

#print(train_data['point'].shape,train_data['feat'].shape,train_data['label'].shape)

print(train_data['point'],train_data['feat'],train_data['label'])

# load valid dataset -> sequences 08

valid_split = dataset.get_split("val")

len = valid_split.__len__()

print('val:',len)

val_data = valid_split.get_data(0)

print(val_data.keys())

#val_data['feat']=np.squeeze(val_data['feat'])

val_data['feat']=None

#print(val_data['point'].shape,val_data['feat'].shape,val_data['label'].shape)

print(val_data['point'],val_data['feat'],val_data['label'])

# load test dataset -> sequences 11~21

test_split = dataset.get_split("test")

len = test_split.__len__()

print('test:',len)

test_data = test_split.get_data(0)

print(test_data.keys())

#test_data['feat']=np.squeeze(test_data['feat'])

test_data['feat']=None

#print(test_data['point'].shape,test_data['feat'].shape,test_data['label'].shape)

print(test_data['point'],test_data['feat'],test_data['label'])

# run inference -> train data중에서 처음 5개 포인트 클라우드 파일에 대하여 추론 수행

for i in range(5):

    train_data = train_split.get_data(i)

    #print(train_data.keys())

    train_data['feat']=None

    #print(train_data['point'].shape,train_data['feat'].shape,train_data['label'].shape)

    #print(train_data['point'],train_data['feat'],train_data['label'])

    pred_custom_data(i, train_data, pipeline_r, pipeline_k)

# visualize prediced results -> 5개 파일에 대하여 2개의 분할모델의 추론 결과를 시각화

v.visualize(vis_points)

3. 실행결과

우측 아래 dataset 창에서 0 -> 원본 데이터, 0_randlanet -> randlanet 모델 추론결과, 0_kpconv -> kpconv 모델 추론결과

원본과 추론결과를 비교해보면 성능을 확인가능

추론 결과 데이터는 포인트 클라우드 데이터의 포인트별로 모델이 예측한 클래스 정보가 추가로 저장되어 있음

시각화 객체에서 클래스정보를 정의한 색상으로 구분하여 출력해줌