MinkowskiEngine语义分割
MinkowskiEngine語義分割
要運(yùn)行示例,請(qǐng)安裝Open3D與PIP安裝open3d-python。
cd /path/to/MinkowskiEngine
python -m examples.indoor
細(xì)分酒店房間
運(yùn)行示例時(shí),將看到一個(gè)旅館房間和房間的語義分割。運(yùn)行示例時(shí),以交互方式旋轉(zhuǎn)可視化效果。
首先,加載數(shù)據(jù)并體素化(量化)數(shù)據(jù)。調(diào)用MinkowskiEngine.utils.sparse_quantize進(jìn)行體素化。
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)
準(zhǔn)備體素化的坐標(biāo)和特征后,應(yīng)用MinkowskiEngine.SparseTensor將其包裹起來。此前,通過調(diào)用MinkowskiEngine.utils.sparse_collate來創(chuàng)建批處理。此函數(shù)采用一組坐標(biāo)和特征并將其連接起來。還將批處理索引附加到坐標(biāo)。最后,通過從顏色中減去0.5,對(duì)特征進(jìn)行偽歸一化。
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)
最后,將稀疏張量前饋到網(wǎng)絡(luò)中并獲得預(yù)測(cè)。
soutput = model(sinput)
_, pred = soutput.F.max(1)
經(jīng)過一些后處理。可以為標(biāo)簽著色,并排可視化輸入和預(yù)測(cè)。
運(yùn)行示例后,權(quán)重會(huì)自動(dòng)下載,并且權(quán)重目前是Scannet 3D分段基準(zhǔn)測(cè)試中排名最高的算法。
有關(guān)更多詳細(xì)信息,請(qǐng)參閱完整的室內(nèi)細(xì)分示例。
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, pcdif __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])
總結(jié)
以上是生活随笔為你收集整理的MinkowskiEngine语义分割的全部?jī)?nèi)容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: MinkowskiEngine demo
- 下一篇: MinkowskiEngine多GPU训