首页 > 其他分享> 文章详细

MinkowskiEngine语义分割

2021-01-04 08:03:11 阅读：379 来源： 互联网

标签：分割 MinkowskiEngine 语义 colors coords file pred import pcd

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])

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

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

ICode9

MinkowskiEngine语义分割