|
- import os
- import numpy as np
- import h5py
-
-
- def read_h5_geo(filedir):
- pc = h5py.File(filedir, 'r')['data'][:]
- coords = pc[:,0:3].astype('int')
-
- return coords
-
- def write_h5_geo(filedir, coords):
- data = coords.astype('uint8')
- with h5py.File(filedir, 'w') as h:
- h.create_dataset('data', data=data, shape=data.shape)
-
- return
-
- def read_ply_ascii_geo(filedir):
- files = open(filedir)
- data = []
- for i, line in enumerate(files):
- wordslist = line.split(' ')
- try:
- line_values = []
- for i, v in enumerate(wordslist):
- if v == '\n': continue
- line_values.append(float(v))
- except ValueError: continue
- data.append(line_values)
- data = np.array(data) #data[0]: [255. 39. 291. 0.505101 0.565322 0.652139 137. 62. 46. 255. ] [x y z nx ny nz r g b alpha]
- coords = data[:,0:3].astype('int')
-
- return coords
-
- def write_ply_ascii_geo(filedir, coords):
- if os.path.exists(filedir): os.system('rm '+filedir)
- f = open(filedir,'a+')
- f.writelines(['ply\n','format ascii 1.0\n'])
- f.write('element vertex '+str(coords.shape[0])+'\n')
- f.writelines(['property float x\n','property float y\n','property float z\n'])
- f.write('end_header\n')
- coords = coords.astype('int')
- for p in coords:
- f.writelines([str(p[0]), ' ', str(p[1]), ' ',str(p[2]), '\n'])
- f.close()
-
- return
-
- ###########################################################################################################
-
- import torch
- import MinkowskiEngine as ME
-
- def array2vector(array, step):
- """ravel 2D array with multi-channel to one 1D vector by sum each channel with different step.
- """
- array, step = array.long().cpu(), step.long().cpu() #array[0]: tensor([0,248,32,288], dtype=torch.int32) step: tensor(969, dtype=torch.int32)
- vector = sum([array[:,i]*(step**i) for i in range(array.shape[-1])]) #array.shape: torch.Size([13849, 4]) vector[0]: tensor(262068011256)
-
- return vector #vector.shape: torch.Size([13849])
-
- def isin(data, ground_truth):
- """ Input data and ground_truth are torch tensor of shape [N, D].
- Returns a boolean vector of the same length as `data` that is True
- where an element of `data` is in `ground_truth` and False otherwise.
- """
- device = data.device
- data, ground_truth = data.cpu(), ground_truth.cpu()
- step = torch.max(data.max(), ground_truth.max()) + 1
- data = array2vector(data, step)
- ground_truth = array2vector(ground_truth, step)
- mask = np.isin(data.cpu().numpy(), ground_truth.cpu().numpy())
-
- return torch.Tensor(mask).bool().to(device)
-
- def istopk(data, nums, rho=1.0):
- """ Input data is sparse tensor and nums is a list of shape [batch_size].
- Returns a boolean vector of the same length as `data` that is True
- where an element of `data` is the top k (=nums*rho) value and False otherwise.
- """
- mask = torch.zeros(len(data), dtype=torch.bool)
- row_indices_per_batch = data._batchwise_row_indices #row_indices_per_batch: [tensor([ 0, 1, 2, ..., 110789, 110790, 110791], device='cuda:0')]
- for row_indices, N in zip(row_indices_per_batch, nums):
- k = int(min(len(row_indices), N*rho))
- _, indices = torch.topk(data.F[row_indices].squeeze().detach().cpu(), k)# must CPU. data.F.shape: torch.Size([110792, 1]),通过F的值判断,最大的k个留,通过num判断留多少个
- mask[row_indices[indices]]=True
-
- return mask.bool().to(data.device)
-
- def sort_spare_tensor(sparse_tensor):
- """ Sort points in sparse tensor according to their coordinates.
- """
- indices_sort = np.argsort(array2vector(sparse_tensor.C.cpu(),
- sparse_tensor.C.cpu().max()+1)) #升序 sparse_tensor.C的shape[13849, 4],在4那个维度中,越是后面的位置对排序结果影响越大
- sparse_tensor_sort = ME.SparseTensor(features=sparse_tensor.F[indices_sort],
- coordinates=sparse_tensor.C[indices_sort],
- tensor_stride=sparse_tensor.tensor_stride[0], #sparse_tensor.tensor_stride: [8, 8, 8]
- device=sparse_tensor.device)
-
- return sparse_tensor_sort
-
- def load_sparse_tensor(filedir, device):
- coords = torch.tensor(read_ply_ascii_geo(filedir)).int()
- feats = torch.ones((len(coords),1)).float()
- # coords, feats = ME.utils.sparse_quantize(coordinates=coords, features=feats, quantization_size=1)
- coords, feats = ME.utils.sparse_collate([coords], [feats])
- x = ME.SparseTensor(features=feats, coordinates=coords, tensor_stride=1, device=device)
-
- return x
|