def knn_search(support_pts, query_pts, k):
"""
:param support_pts: points you have, B*N1*3
:param query_pts: points you want to know the neighbour index, B*N2*3
:param k: Number of neighbours in knn search
:return: neighbor_idx: neighboring points indexes, B*N2*k
"""
neighbor_idx = nearest_neighbors.knn_batch(support_pts,
query_pts,
k,
omp=True)
return neighbor_idx.astype(np.int32)
def nearest_neighbor_interpolate(unknown, known, known_feats):
"""
:param pts: (n, 4) tensor of the bxyz positions of the unknown features
:param ctr: (m, 4) tensor of the bxyz positions of the known features
:param ctr_feats: (m, C) sparsetensor of features to be propigated
:return:
new_features: (n, C) tensor of the features of the unknown features
"""
batch_size = known_feats.batch_size
features = known_feats.features
channel = features.shape[-1]
ind = torch.zeros(unknown.shape[0], 1)
offsetq = 0
offsets = 0
for b in range(batch_size):
support_points = known[known[:, 0] == b][:, 1:].unsqueeze(0).to(
'cpu').data.numpy() # [N,3]
query_points = unknown[unknown[:, 0] == b][:, 1:].unsqueeze(0).to(
'cpu').data.numpy() # [M,3]
indexs = nearest_neighbors.knn_batch(support_points,
query_points,
1,
omp=True)
num_spoint = support_points.shape[1]
num_qpoint = query_points.shape[1]
ind[offsetq:num_qpoint + offsetq] = torch.Tensor(indexs +
offsets).squeeze(0)
offsets += num_spoint
offsetq += num_qpoint
ind = ind.long()
interpolated_feat = features[ind.long()].view(-1, channel)
return interpolated_feat.contiguous()
def collate_batch(self, batch_list, _unused=False):
data_dict = defaultdict(list)
for cur_sample in batch_list:
for key, val in cur_sample.items():
data_dict[key].append(val)
batch_size = len(batch_list)
ret = {}
image_idx = data_dict['frame_id']
calib = data_dict['calib']
for key, val in data_dict.items():
try:
if key in ['image']:
ret[key] = np.concatenate(val, axis=0)
elif key in ['points']:
num = []
for i, coor in enumerate(val):
num.append(coor.shape[0])
## sample raw_points & voxels_feature
num_ = np.max(num)
coors_batch = []
#if num_ > 10000:
# num_ = 10000
for i, coor in enumerate(val):
p, _ = sample_points(coor, coor, num_)
coors_batch.append(p)
k = key + '_batch'
ret[k] = np.array(coors_batch).reshape(
(len(num), num_, -1))
elif key in ['raw_points']:
if self.raw == False: # dont need load raw data
#print(self.nbg)
num = []
for i, coor in enumerate(val):
num.append(coor.shape[0])
## sample raw_points
num_ = np.max(num)
coors_batch = []
#if num_ > 10000:
# num_ = 10000
for i, coor in enumerate(val):
p, _ = sample_points(coor, coor, num_)
if self.use_color:
img_path = '/media/ddd/data2/3d_MOTS_Ex./Code/OpenPCDet-RandlaNet/data/kitti/training/image_2/' + \
image_idx[i] + '.png'
img = cv2.imread(img_path)
image = np.float32(img)
if not self.use_rgb:
img = cv2.cvtColor(image,
cv2.COLOR_BGR2Lab)
p = p[:, :3]
calib_result = calib[i].lidar_to_img(
p) # [N,3] in lidar to [N,2] in img
p = self.painted_point_cloud(
calib_result, img, p)
coors_batch.append(p)
k = key + '_batch'
ret[k] = np.array(coors_batch).reshape(
(len(num), num_,
-1)) # if colored,return have colored point cloud
elif key in ['gt_boxes']:
max_gt = max([len(x) for x in val])
batch_gt_boxes3d = np.zeros(
(batch_size, max_gt, val[0].shape[-1]),
dtype=np.float32)
for k in range(batch_size):
batch_gt_boxes3d[k, :val[k].__len__(), :] = val[k]
ret[key] = batch_gt_boxes3d
elif key in [
'voxel_features', 'voxels', 'voxel_num_points',
'voxel_coords'
]:
continue
else:
ret[key] = np.stack(val, axis=0)
except:
print('image_idx:', image_idx)
img_path = '/media/ddd/data2/3d_MOTS_Ex./Code/OpenPCDet-RandlaNet/data/kitti/training/image_2/' + \
image_idx[0] + '.png'
img = cv2.imread(img_path)
print(img_path)
print('Error in collate_batch: key=%s' % key)
raise TypeError
val = data_dict['voxel_features']
voxel_coords = data_dict['voxel_coords']
num = []
for i, coor in enumerate(val):
num.append(coor.shape[0])
num_ = np.max(num)
fea = []
cor = []
#if num_ > 10000:
# num_ = 10000
for i, coor in enumerate(val):
p, c = sample_points(coor, voxel_coords[i], num_)
if self.use_color:
img_path = '/media/ddd/data2/3d_MOTS_Ex./Code/OpenPCDet-RandlaNet/data/kitti/training/image_2/' + \
image_idx[i] + '.png'
img = cv2.imread(img_path)
image = np.float32(img)
if not self.use_rgb:
img = cv2.cvtColor(image, cv2.COLOR_BGR2Lab)
p = p[:, :3]
calib_result = calib[i].lidar_to_img(
p) # [N,3] in lidar to [N,2] in img
p = self.painted_point_cloud(calib_result, img, p)
fea.append(p)
cor.append(c)
ret['voxel_features_batch'] = np.array(fea).reshape(
(len(num), num_, -1)) #if colored,return have colored point cloud
ret['voxel_coords_batch'] = np.array(cor).reshape((len(num), num_, -1))
val = ret['voxel_coords_batch']
coors = []
for i, coor in enumerate(val):
coor_pad = np.pad(coor, ((0, 0), (1, 0)),
mode='constant',
constant_values=i)
coors.append(coor_pad)
ret['voxel_coords_batch'] = np.concatenate(coors, axis=0)
voxel_features_batch = ret['voxel_features_batch'][:, :, :3]
if self.nbg: # True ,使用有背景点的选点方式
raw_points_batch = ret['raw_points_batch'][:, :, :3]
neighbor_idx = nearest_neighbors.knn_batch(raw_points_batch,
voxel_features_batch,
8,
omp=True)
else:
neighbor_idx = nearest_neighbors.knn_batch(voxel_features_batch,
voxel_features_batch,
8,
omp=True)
ret['neighbor'] = neighbor_idx
ret['batch_size'] = batch_size
return ret
def forward(self, feat1, feat2):
'''feat1 is sparsetensor and feat2 is densetensor'''
device = feat1.features.device
spatial_shape = feat2.spatial_shape
indices = feat2.indices
batch_size = feat1.batch_size
coord1, features1 = extract_coord_features(feat1.dense().detach())
coord2, features2 = extract_coord_features(feat2.dense())
xyz1 = align_pnt(coord1, self.voxelsize,
self.point_range) # [0, -25.6, -2, 51.2, 25.6, 4.4]
xyz2 = align_pnt(coord2, self.voxelsize * 2, self.point_range)
ind = torch.zeros(features2.shape[0], self.k)
offsets = 0
offsetq = 0
for b in range(batch_size):
support_points = xyz1[xyz1[:, 0] == b][:, 1:].unsqueeze(0).to(
'cpu').data.numpy() # [N,3]
query_points = xyz2[xyz2[:, 0] == b][:, 1:].unsqueeze(0).to(
'cpu').data.numpy() # [M,3]
indexs = nearest_neighbors.knn_batch(support_points,
query_points,
self.k,
omp=True)
assert len(indexs[0]) == len(query_points[0])
num_spoint = support_points.shape[1]
num_qpoint = query_points.shape[1]
ind[offsetq:num_qpoint +
offsetq] = torch.Tensor(indexs + offsets).squeeze(0)
offsets += num_spoint
offsetq += num_qpoint
ind = ind.long().to(device)
group, _ = VoxelPooling.index_feat(
torch.cat([xyz1[:, 1:].float(), features1], 1), ind)
group_xyz, group_features = group[:, :, :
3], group[:, :,
3:] # [N, K, 3], [N, K, D]
relation = VoxelPooling.relation_position(
group_xyz, xyz2[:, 1:].float()) # [N, K, 10]
if self.feat_relation:
relation = torch.cat([
relation, group_features,
features2.unsqueeze(1).repeat((1, self.k, 1))
], -1)
group_w = self.relation_w(relation.permute(0, 2, 1))
group_features = group_features.permute(0, 2, 1)
group_features *= group_w
updated_features = torch.mean(group_features, 2)
features2 = self.fuse_mlps(torch.cat([updated_features, features2], 1))
feat = spconv.SparseConvTensor(features2.float(), indices.int(),
spatial_shape, batch_size)
feat.indice_dict = feat2.indice_dict
feat.grid = feat2.grid
return feat.dense()
import numpy as np #import lib.python.nearest_neighbors as nearest_neighbors import nearest_neighbors import time batch_size = 16 num_points = 81920 K = 16 pc = np.random.rand(batch_size, num_points, 3).astype(np.float32) # nearest neighbours start = time.time() neigh_idx = nearest_neighbors.knn_batch(pc, pc, K, omp=True) print(time.time() - start)