其他分享
首页 > 其他分享> > MinkowskiEngine语义分割

MinkowskiEngine语义分割

作者:互联网

MinkowskiEngine语义分割

要运行示例,请安装Open3DPIP安装open3d-python

cd /path/to/MinkowskiEngine

python -m examples.indoor

细分酒店房间

运行示例时,将看到一个旅馆房间和房间的语义分割。运行示例时,以交互方式旋转可视化效果。

首先,加载数据并体素化(量化)数据。调用MinkowskiEngine.utils.sparse_quantize进行体素化。

pcd = o3d.read_point_cloud(file_name)

coords = np.array(pcd.points)

feats = np.array(pcd.colors)

 

quantized_coords = np.floor(coords / voxel_size)

inds = ME.utils.sparse_quantize(quantized_coords)

准备体素化的坐标和特征后,应用MinkowskiEngine.SparseTensor将其包裹起来。此前,通过调用MinkowskiEngine.utils.sparse_collate来创建批处理。此函数采用一组坐标和特征并将其连接起来。还将批处理索引附加到坐标。最后,通过从颜色中减去0.5,对特征进行伪归一化。

# Create a batch, this process is done in a data loader during training in parallel.

batch = [load_file(config.file_name, 0.02)]

coordinates_, featrues_, pcds = list(zip(*batch))

coordinates, features = ME.utils.sparse_collate(coordinates_, featrues_)

 

# Normalize features and create a sparse tensor

sinput = ME.SparseTensor(features - 0.5, coords=coordinates).to(device)

最后,将稀疏张量前馈到网络中并获得预测。

soutput = model(sinput)

_, pred = soutput.F.max(1)

经过一些后处理。可以为标签着色,并排可视化输入和预测。

 

 

 运行示例后,权重会自动下载,并且权重目前是Scannet 3D分段基准测试中排名最高的算法。

有关更多详细信息,请参阅完整的室内细分示例

import os

 

import argparse

 

import numpy as np

 

from urllib.request import urlretrieve

 

try:

 

import open3d as o3d

 

except ImportError:

 

raise ImportError('Please install open3d with `pip install open3d`.')

   
 

import torch

 

import MinkowskiEngine as ME

 

from examples.minkunet import MinkUNet34C

 

from examples.common import Timer

   
 

# Check if the weights and file exist and download

 

if not os.path.isfile('weights.pth'):

 

print('Downloading weights and a room ply file...')

 

urlretrieve("http://cvgl.stanford.edu/data2/minkowskiengine/weights.pth",

 

'weights.pth')

 

urlretrieve("http://cvgl.stanford.edu/data2/minkowskiengine/1.ply", '1.ply')

   
 

parser = argparse.ArgumentParser()

 

parser.add_argument('--file_name', type=str, default='1.ply')

 

parser.add_argument('--weights', type=str, default='weights.pth')

 

parser.add_argument('--use_cpu', action='store_true')

   
 

CLASS_LABELS = ('wall', 'floor', 'cabinet', 'bed', 'chair', 'sofa', 'table',

 

'door', 'window', 'bookshelf', 'picture', 'counter', 'desk',

 

'curtain', 'refrigerator', 'shower curtain', 'toilet', 'sink',

 

'bathtub', 'otherfurniture')

   
 

VALID_CLASS_IDS = [

 

1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 16, 24, 28, 33, 34, 36, 39

 

]

   
 

SCANNET_COLOR_MAP = {

 

0: (0., 0., 0.),

 

1: (174., 199., 232.),

 

2: (152., 223., 138.),

 

3: (31., 119., 180.),

 

4: (255., 187., 120.),

 

5: (188., 189., 34.),

 

6: (140., 86., 75.),

 

7: (255., 152., 150.),

 

8: (214., 39., 40.),

 

9: (197., 176., 213.),

 

10: (148., 103., 189.),

 

11: (196., 156., 148.),

 

12: (23., 190., 207.),

 

14: (247., 182., 210.),

 

15: (66., 188., 102.),

 

16: (219., 219., 141.),

 

17: (140., 57., 197.),

 

18: (202., 185., 52.),

 

19: (51., 176., 203.),

 

20: (200., 54., 131.),

 

21: (92., 193., 61.),

 

22: (78., 71., 183.),

 

23: (172., 114., 82.),

 

24: (255., 127., 14.),

 

25: (91., 163., 138.),

 

26: (153., 98., 156.),

 

27: (140., 153., 101.),

 

28: (158., 218., 229.),

 

29: (100., 125., 154.),

 

30: (178., 127., 135.),

 

32: (146., 111., 194.),

 

33: (44., 160., 44.),

 

34: (112., 128., 144.),

 

35: (96., 207., 209.),

 

36: (227., 119., 194.),

 

37: (213., 92., 176.),

 

38: (94., 106., 211.),

 

39: (82., 84., 163.),

 

40: (100., 85., 144.),

 

}

   
   
 

def load_file(file_name):

 

pcd = o3d.io.read_point_cloud(file_name)

 

coords = np.array(pcd.points)

 

colors = np.array(pcd.colors)

 

return coords, colors, pcd

   
   
 

if __name__ == '__main__':

 

config = parser.parse_args()

 

device = torch.device('cuda' if (

 

torch.cuda.is_available() and not config.use_cpu) else 'cpu')

 

print(f"Using {device}")

 

# Define a model and load the weights

 

model = MinkUNet34C(3, 20).to(device)

 

model_dict = torch.load(config.weights)

 

model.load_state_dict(model_dict)

 

model.eval()

   
 

coords, colors, pcd = load_file(config.file_name)

 

# Measure time

 

with torch.no_grad():

 

voxel_size = 0.02

 

# Feed-forward pass and get the prediction

 

in_field = ME.TensorField(

 

features=torch.from_numpy(colors).float(),

 

coordinates=ME.utils.batched_coordinates([coords / voxel_size], dtype=torch.float32),

 

quantization_mode=ME.SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE,

 

minkowski_algorithm=ME.MinkowskiAlgorithm.SPEED_OPTIMIZED,

 

device=device,

 

)

 

# Convert to a sparse tensor

 

sinput = in_field.sparse()

 

# Output sparse tensor

 

soutput = model(sinput)

 

# get the prediction on the input tensor field

 

out_field = soutput.slice(in_field)

 

logits = out_field.F

   
 

_, pred = logits.max(1)

 

pred = pred.cpu().numpy()

   
 

# Create a point cloud file

 

pred_pcd = o3d.geometry.PointCloud()

 

# Map color

 

colors = np.array([SCANNET_COLOR_MAP[VALID_CLASS_IDS[l]] for l in pred])

 

pred_pcd.points = o3d.utility.Vector3dVector(coords)

 

pred_pcd.colors = o3d.utility.Vector3dVector(colors / 255)

 

pred_pcd.estimate_normals()

   
 

# Move the original point cloud

 

pcd.points = o3d.utility.Vector3dVector(

 

np.array(pcd.points) + np.array([0, 5, 0]))

   
 

# Visualize the input point cloud and the prediction

 

o3d.visualization.draw_geometries([pcd, pred_pcd])

标签:分割,MinkowskiEngine,语义,colors,coords,file,pred,import,pcd
来源: https://www.cnblogs.com/wujianming-110117/p/14227739.html