MinkowskiEngine语义分割
要运行示例，请安装Open3D与PIP安装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]) |