def _infer_flow(self, flow_f_1, flow_f_2): # Normalize the features if self.normalize_feature: flow_f_1 = ME.SparseTensor( flow_f_1.F / torch.norm(flow_f_1.F, p=2, dim=1, keepdim=True), coordinate_map_key=flow_f_1.coordinate_map_key, coordinate_manager=flow_f_1.coordinate_manager) flow_f_2 = ME.SparseTensor( flow_f_2.F / torch.norm(flow_f_2.F, p=2, dim=1, keepdim=True), coordinate_map_key=flow_f_2.coordinate_map_key, coordinate_manager=flow_f_2.coordinate_manager) # Extract the coarse flow based on the feature correspondences coarse_flow = [] # Iterate over the examples in the batch for b_idx in range(len(flow_f_1.decomposed_coordinates)): feat_s = flow_f_1.F[flow_f_1.C[:, 0] == b_idx] feat_t = flow_f_2.F[flow_f_2.C[:, 0] == b_idx] coor_s = flow_f_1.C[flow_f_1.C[:, 0] == b_idx, 1:].to( self.device) * self.voxel_size coor_t = flow_f_2.C[flow_f_2.C[:, 0] == b_idx, 1:].to( self.device) * self.voxel_size # Squared l2 distance between points points of both point clouds coor_s, coor_t = coor_s.unsqueeze(0), coor_t.unsqueeze(0) feat_s, feat_t = feat_s.unsqueeze(0), feat_t.unsqueeze(0) # Force transport to be zero for points further than 10 m apart support = (pairwise_distance(coor_s, coor_t, normalized=False) < 10 **2).float() # Transport cost matrix C = pairwise_distance(feat_s, feat_t) K = torch.exp( -C / (torch.exp(self.epsilon) + self.tau_offset)) * support row_sum = K.sum(-1, keepdim=True) # Estimate flow corr_flow = (K @ coor_t) / (row_sum + 1e-8) - coor_s coarse_flow.append(corr_flow.squeeze(0)) coarse_flow = torch.cat(coarse_flow, dim=0) st_cf = ME.SparseTensor(features=coarse_flow, coordinate_manager=flow_f_1.coordinate_manager, coordinate_map_key=flow_f_1.coordinate_map_key) self.inferred_values['coarse_flow'] = st_cf.F # Refine the flow with the second network refined_flow = self.flow_refiner(st_cf) self.inferred_values['refined_flow'] = refined_flow.F
def pose_estimation(model, device, xyz0, xyz1, coord0, coord1, feats0, feats1, return_corr=False): sinput0 = ME.SparseTensor(feats0.to(device), coordinates=coord0.to(device)) F0 = model(sinput0).F sinput1 = ME.SparseTensor(feats1.to(device), coordinates=coord1.to(device)) F1 = model(sinput1).F corr = F0.mm(F1.t()) weight, inds = corr.max(dim=1) weight = weight.unsqueeze(1).cpu() xyz1_corr = xyz1[inds, :] trans = est_quad_linear_robust(xyz0, xyz1_corr, weight) # let's do this later if return_corr: return trans, weight, corr else: return trans, weight
def empty_return(self): device = self.efi_cls.kernel.device dtype = self.efi_cls.kernel.dtype gfi = ME.SparseTensor( coordinates=torch.empty((0, self.spatial_dims + 1), device=device), features=torch.empty((0, self.efi.out_connection_type[1]), device=device, dtype=dtype), tensor_stride=[2**(self.config.n_conv_layers)] * self.spatial_dims) lfi = ME.SparseTensor( coordinates=torch.empty((0, self.spatial_dims + 1), device=device), features=torch.empty((0, self.gfi.out_connection_type[1]), device=device, dtype=dtype), tensor_stride=[2**(self.config.feature_layer + 1)] * self.spatial_dims) recon_x = ME.SparseTensor(coordinates=torch.empty( (0, self.spatial_dims + 1), device=device), features=torch.empty( (0, self.lfi.out_connection_type[1]), device=device, dtype=dtype), tensor_stride=[1] * self.spatial_dims) empty_outputs = {"gfi": gfi, "lfi": lfi, "recon_x": recon_x} return empty_outputs
def _compute_loss_metrics(self, input_dict, phase='train'): ''' Computes the losses and evaluation metrics Args: input_dict (dict): data dictionary Return: losses (dict): selected loss values metric (dict): selected evaluation metric ''' # Run the feature and context encoder sinput1 = ME.SparseTensor(features=input_dict['sinput_s_F'].to(self.device), coordinates=input_dict['sinput_s_C'].to(self.device)) sinput2 = ME.SparseTensor(features=input_dict['sinput_t_F'].to(self.device), coordinates=input_dict['sinput_t_C'].to(self.device)) if phase == 'train': inferred_values = self.model(sinput1, sinput2, input_dict['pcd_s'], input_dict['pcd_t'], input_dict['fg_labels_s'], input_dict['fg_labels_t']) else: inferred_values = self.model(sinput1, sinput2, input_dict['pcd_eval_s'], input_dict['pcd_eval_t'], input_dict['fg_labels_s'], input_dict['fg_labels_t']) losses = self.compute_losses(inferred_values, input_dict) metrics = self.compute_metrics(inferred_values, input_dict, phase) return losses, metrics
def forward(self, x): xf = self.relu( self.bn1(self.conv1(x.F.unsqueeze(-1))[..., 0]).unsqueeze(-1)) xf = self.relu(self.bn2(self.conv2(xf)[..., 0]).unsqueeze(-1)) xf = self.relu(self.bn3(self.conv3(xf)[..., 0]).unsqueeze(-1)) xf = ME.SparseTensor(xf[..., 0], coords_key=x.coords_key, coords_manager=x.coords_man) xfc = self.pool(xf) xf = F.relu(self.bn4(self.fc1(self.pool(xfc).F))) xf = F.relu(self.bn5(self.fc2(xf))) xf = self.fc3(xf) xf += torch.tensor([[1, 0, 0, 0, 1, 0, 0, 0, 1]], dtype=x.dtype, device=x.device).repeat(xf.shape[0], 1) xf = ME.SparseTensor(xf, coords_key=xfc.coords_key, coords_manager=xfc.coords_man) xfc = ME.SparseTensor(torch.zeros(x.shape[0], 9, dtype=x.dtype, device=x.device), coords_key=x.coords_key, coords_manager=x.coords_man) return self.broadcast(xfc, xf)
def benchmark(config): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = ResUNetBN2C(1, 16, normalize_feature=True, conv1_kernel_size=3, D=3) model.eval() model = model.to(device) num_conv_layers = defaultdict(int) for l in model.modules(): if isinstance(l, ME.MinkowskiConvolution) or isinstance( l, ME.MinkowskiConvolutionTranspose ): num_conv_layers[l.kernel_generator.kernel_size[0]] += 1 print(num_conv_layers) pcd = o3d.io.read_point_cloud(config.input) if ME.__version__.split(".")[1] == "5": for batch_size in [1, 2, 4, 8, 16, 32, 64]: vox_coords = torch.from_numpy(np.array(pcd.points)) / config.voxel_size coords = ME.utils.batched_coordinates( [vox_coords for i in range(batch_size)] ) feats = torch.from_numpy(np.ones((len(coords), 1))).float() with torch.no_grad(): t = Timer() for i in range(10): # initialization time includes copy to GPU t.tic() sinput = ME.SparseTensor( feats, coords, minkowski_algorithm=ME.MinkowskiAlgorithm.SPEED_OPTIMIZED, device=device, ) model(sinput) t.toc() print(f"{batch_size}\t{len(sinput)}\t{t.min_time}") elif ME.__version__.split(".")[1] == "4": for batch_size in [1, 2, 4, 8, 16, 32, 64]: vox_coords = torch.from_numpy(np.array(pcd.points)) / config.voxel_size coords = ME.utils.batched_coordinates( [vox_coords for i in range(batch_size)] ) feats = torch.from_numpy(np.ones((len(coords), 1))).float() with torch.no_grad(): t = Timer() for i in range(10): # initialization time includes copy to GPU t.tic() sinput = ME.SparseTensor(feats, coords,).to(device) model(sinput) t.toc() print(f"{batch_size}\t{len(sinput)}\t{t.min_time}") else: raise NotImplementedError
def corr_and_add(feature_A, feature_B, k=10, coords_A=None, coords_B=None, Npts=None): # compute sparse correlation from A to B scorr = sparse_corr(feature_A, feature_B, k=k, ratio=False, sparse_type='raw') # compute sparse correlation from B to A scorr2 = sparse_corr(feature_B, feature_A, k=k, ratio=False, reverse=True, sparse_type='raw') scorr = ME.SparseTensor(scorr[0], scorr[1].cuda()) scorr2 = ME.SparseTensor(scorr2[0], scorr2[1].cuda(), coordinate_manager=scorr.coordinate_manager) scorr = ME.MinkowskiUnion()(scorr, scorr2) return scorr
def generate_inlier_input(self, xyz0, xyz1, iC0, iC1, iF0, iF1, len_batch, pos_pairs): # pairs consist of (xyz1 index, xyz0 index) stime = time.time() sinput0 = ME.SparseTensor(iF0, coords=iC0).to(self.device) oF0 = self.feat_model(sinput0).F sinput1 = ME.SparseTensor(iF1, coords=iC1).to(self.device) oF1 = self.feat_model(sinput1).F feat_time = time.time() - stime stime = time.time() pred_pairs = self.find_pairs(oF0, oF1, len_batch) nn_time = time.time() - stime is_correct = find_correct_correspondence(pos_pairs, pred_pairs, len_batch=len_batch) cat_pred_pairs = [] start_inds = torch.zeros((1, 2)).long() for lens, pred_pair in zip(len_batch, pred_pairs): cat_pred_pairs.append(pred_pair + start_inds) start_inds += torch.LongTensor(lens) cat_pred_pairs = torch.cat(cat_pred_pairs, 0) pred_pair_inds0, pred_pair_inds1 = cat_pred_pairs.t() reg_coords = torch.cat( (iC0[pred_pair_inds0], iC1[pred_pair_inds1, 1:]), 1) reg_feats = self.generate_inlier_features(xyz0, xyz1, iC0, iC1, oF0, oF1, pred_pair_inds0, pred_pair_inds1).float() return reg_coords, reg_feats, pred_pairs, is_correct, feat_time, nn_time
def forward(self, x): # First, align coordinates to be centered around zero. coords = x.coords.to(x.device)[:, 1:] coords = ME.SparseTensor(coords.float(), coords_key=x.coords_key, coords_manager=x.coords_man) mean_coords = self.broadcast(coords, self.pool(coords)) norm_coords = coords - mean_coords # Second, apply spatial transformer to the coordinates. trans = self.stn(norm_coords) coords_feat_stn = torch.squeeze( torch.bmm(norm_coords.F.view(-1, 1, 3), trans.F.view(-1, 3, 3))) xf = torch.cat((coords_feat_stn, x.F), 1).unsqueeze(-1) xf = self.relu(self.bn1(self.conv1(xf)[..., 0]).unsqueeze(-1)) pointfeat = xf xf = self.bn2(self.conv2(xf)[..., 0]).unsqueeze(-1) xfc = ME.SparseTensor(xf[..., 0], coords_key=x.coords_key, coords_manager=x.coords_man) xf_avg = ME.SparseTensor(torch.zeros(x.shape[0], xfc.F.shape[1], dtype=x.dtype, device=x.device), coords_key=x.coords_key, coords_manager=x.coords_man) xf_avg = self.broadcast(xf_avg, self.pool(xfc)) return torch.cat((pointfeat[..., 0], xf_avg.F), 1)
def operation_mode(): # Set to share the coords_man by default ME.set_sparse_tensor_operation_mode( ME.SparseTensorOperationMode.SHARE_COORDS_MANAGER) print(ME.sparse_tensor_operation_mode()) coords0, feats0 = to_sparse_coo(data_batch_0) coords0, feats0 = ME.utils.sparse_collate(coords=[coords0], feats=[feats0]) coords1, feats1 = to_sparse_coo(data_batch_1) coords1, feats1 = ME.utils.sparse_collate(coords=[coords1], feats=[feats1]) for _ in range(2): # sparse tensors A = ME.SparseTensor(coords=coords0, feats=feats0) B = ME.SparseTensor( coords=coords1, feats=feats1, # coords_manager=A.coords_man, No need to feed the coords_man force_creation=True) C = A + B # When done using it for forward and backward, you must cleanup the coords man ME.clear_global_coords_man()
def forward(self, x): out_s1 = self.conv1(x) out_s1 = self.norm1(out_s1) out_s1 = self.block1(out_s1) out = MEF.relu(out_s1) out_s2 = self.conv2(out) out_s2 = self.norm2(out_s2) out_s2 = self.block2(out_s2) out = MEF.relu(out_s2) out_s4 = self.conv3(out) out_s4 = self.norm3(out_s4) out_s4 = self.block3(out_s4) out = MEF.relu(out_s4) out_s8 = self.conv4(out) out_s8 = self.norm4(out_s8) out_s8 = self.block4(out_s8) out = MEF.relu(out_s8) out = self.conv4_tr(out) out = self.norm4_tr(out) out = self.block4_tr(out) out_s4_tr = MEF.relu(out) out = ME.cat(out_s4_tr, out_s4) out = self.conv3_tr(out) out = self.norm3_tr(out) out = self.block3_tr(out) out_s2_tr = MEF.relu(out) out = ME.cat(out_s2_tr, out_s2) out = self.conv2_tr(out) out = self.norm2_tr(out) out = self.block2_tr(out) out_s1_tr = MEF.relu(out) out = ME.cat(out_s1_tr, out_s1) out_feat = self.conv1_tr(out) out_feat = MEF.relu(out_feat) out_feat = self.final(out_feat) out_att = self.conv1_tr_att(out) out_att = MEF.relu(out_att) out_att = self.final_att(out_att) out_att = ME.SparseTensor(self.final_att_act(out_att.F), coords_key=out_att.coords_key, coords_manager=out_att.coords_man) if self.normalize_feature: out_feat = ME.SparseTensor( out_feat.F / torch.norm(out_feat.F, p=2, dim=1, keepdim=True), coords_key=out_feat.coords_key, coords_manager=out_feat.coords_man) return dict(features=out_feat, attention=out_att)
def forward(self, x): out_s1 = self.conv1(x) out_s1 = self.norm1(out_s1) out_s1 = self.block1(out_s1) out = MEF.relu(out_s1) out_s2 = self.conv2(out) out_s2 = self.norm2(out_s2) out_s2 = self.block2(out_s2) out = MEF.relu(out_s2) out_s4 = self.conv3(out) out_s4 = self.norm3(out_s4) out_s4 = self.block3(out_s4) out = MEF.relu(out_s4) out_s8 = self.conv4(out) out_s8 = self.norm4(out_s8) out_s8 = self.block4(out_s8) out = MEF.relu(out_s8) out = self.conv4_tr(out) out = self.norm4_tr(out) out = self.block4_tr(out) out_s4_tr = MEF.relu(out) out = ME.cat(out_s4_tr, out_s4) out = self.conv3_tr(out) out = self.norm3_tr(out) out = self.block3_tr(out) out_s2_tr = MEF.relu(out) out = ME.cat(out_s2_tr, out_s2) out = self.conv2_tr(out) out = self.norm2_tr(out) out = self.block2_tr(out) out_s1_tr = MEF.relu(out) out = ME.cat(out_s1_tr, out_s1) out = self.conv1_tr(out) out = MEF.relu(out) out = self.final(out) if self.normalize_feature: if ME.__version__.split(".")[1] == "5": return ME.SparseTensor( out.F / torch.norm(out.F, p=2, dim=1, keepdim=True), coordinate_map_key=out.coordinate_map_key, coordinate_manager=out.coordinate_manager) elif ME.__version__.split(".")[1] == "4": return ME.SparseTensor( out.F / torch.norm(out.F, p=2, dim=1, keepdim=True), coords_key=out.coords_key, coords_manager=out.coords_man) else: return out
def compute_descriptors(self, input_dict): ''' If not precomputed it infers the feature descriptors and returns the established correspondences together with the ground truth transformation parameters and inlier labels. Args: input_dict (dict): input data ''' if not self.precomputed_desc: xyz_down = input_dict['sinput0_C'] sinput0 = ME.SparseTensor( input_dict['sinput0_F'], coords=input_dict['sinput0_C']).to(self.device) F0 = self.descriptor_module(sinput0).F # If the FCGF descriptor should be trained with the FCGF loss (need also corresponding desc.) if self.train_descriptor: sinput1 = ME.SparseTensor( input_dict['sinput1_F'], coords=input_dict['sinput1_C']).to(self.device) F1 = self.descriptor_module(sinput1).F else: F1 = torch.empty(F0.shape[0], 0).to(self.device) # Sample the points xyz_batch, f_batch = self.sampler(xyz_down, F0, input_dict['pts_list']) # Build point cloud pairs for the inference xyz_s, xyz_t, f_s, f_t = extract_overlaping_pairs(xyz_batch, f_batch, self.connectivity_info) # Compute nearest neighbors in feature space nn_C_s_t = self.feature_matching(f_s, f_t, xyz_t) # NNs of the source points in the target point cloud nn_C_t_s = self.feature_matching(f_t, f_s, xyz_s) # NNs of the target points in the source point cloud if self.mutuals_flag: mutuals = extract_mutuals(xyz_s, xyz_t, nn_C_s_t, nn_C_t_s) else: mutuals = None # Prepare the input for the filtering block filtering_input = construct_filtering_input_data(xyz_s, xyz_t, input_dict, self.mutuals) else: filtering_input = input_dict F0 = None F1 = None return filtering_input, F0, F1
def _train_iter(self, data_loader_iter, timers): self.model.train() data_meter, data_timer, total_timer = timers self.optimizer.zero_grad() batch_pos_loss, batch_neg_loss, batch_loss = 0, 0, 0 data_time = 0 total_timer.tic() data_timer.tic() input_dict = data_loader_iter.next() data_time += data_timer.toc(average=False) sinput0 = ME.SparseTensor(input_dict['sinput0_F'], coords=input_dict['sinput0_C']).to( self.cur_device) F0 = self.model(sinput0).F sinput1 = ME.SparseTensor(input_dict['sinput1_F'], coords=input_dict['sinput1_C']).to( self.cur_device) F1 = self.model(sinput1).F pos_pairs = input_dict['correspondences'] pos_loss, neg_loss = self.contrastive_hardest_negative_loss( F0, F1, pos_pairs, num_pos=self.config.trainer.num_pos_per_batch * self.batch_size, num_hn_samples=self.config.trainer.num_hn_samples_per_batch * self.batch_size) loss = pos_loss + neg_loss loss.backward() result = {"loss": loss, "pos_loss": pos_loss, "neg_loss": neg_loss} if self.config.misc.num_gpus > 1: result = du.scaled_all_reduce_dict(result, self.config.misc.num_gpus) batch_loss += result["loss"].item() batch_pos_loss += result["pos_loss"].item() batch_neg_loss += result["neg_loss"].item() self.optimizer.step() torch.cuda.empty_cache() total_timer.toc() data_meter.update(data_time) return batch_loss, batch_pos_loss, batch_neg_loss
def subsample_aux(x0, x1, aux, kernel_size=2): aux = ME.SparseTensor(coordinates=aux.C, features=(aux.F + 1)) n_ks = kernel_size**3 coords = x1.C batch_id = coords[:, 0] batch_id = batch_id.unsqueeze(-1).repeat(1, n_ks).reshape(-1, 1) pooled_C = coords[:, 1:] kernel_C = pooled_C.unsqueeze(1).repeat(1, n_ks, 1) diffs = torch.tensor([ [0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1], ], device='cuda') kernel_C = kernel_C + diffs query_C = torch.cat([batch_id, kernel_C.reshape(-1, 3)], dim=1).float() query_F = aux.features_at_coordinates(query_C).reshape(-1, 8) # [N, 8] # find the most freq: find normal will always return 0 # out, _ = torch.mode(query_F) freqs = torch.stack([(query_F == i).sum(dim=1) for i in range(1, 21)]) out = freqs.max(dim=0)[1] # shape [N] out = out - 1 # [-1~20] # debug dict for plot # d = {} # d['origin_pc'] = x0.C # d['origin_pred'] = aux.F # d['new_pc'] = x1.C # d['new_pred'] = out # torch.save(d, './aux.pth') out = ME.SparseTensor(coordinates=x1.C, features=out.float().reshape([-1, 1]).cuda()) return out
def torch_to_me(sten): sten = sten.coalesce() indices = sten.indices().t().contiguous().int() #.cpu() features = sten.values() if len(features.shape) == 1: features = features.unsqueeze(1) return ME.SparseTensor(features, indices)
def benchmark(config): device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = ResUNetBN2C(1, 16, normalize_feature=True, conv1_kernel_size=3, D=3) model.eval() model = model.to(device) pcd = o3d.io.read_point_cloud(config.input) coords = ME.utils.batched_coordinates( [torch.from_numpy(np.array(pcd.points)) / config.voxel_size]) feats = torch.from_numpy(np.ones((len(coords), 1))).float() with torch.no_grad(): t = MinTimer() for i in range(100): # initialization time includes copy to GPU t.tic() sinput = ME.SparseTensor( feats, coords, minkowski_algorithm=ME.MinkowskiAlgorithm.SPEED_OPTIMIZED, # minkowski_algorithm=ME.MinkowskiAlgorithm.MEMORY_EFFICIENT, device=device) model(sinput) t.toc() print(t.min)
def reform_input(coords, labels, device): labels = np.int32(labels) feats = coords coords, feats, labels = coords[0], feats[0], labels[0] if config['data']['quantize']: quantization_size = 0.005 #* 10**((config['num_point'] / 131072) - 1) # = 0.005 for num_point==131072, 0.05 for num_point== 262144 coords, feats, labels = ME.utils.sparse_quantize( coords=coords, feats=feats, labels=labels, quantization_size=quantization_size) if config['data']['batch_coords']: coords_pt = ME.utils.batched_coordinates([coords * 1e6]) else: coords = np.concatenate((coords * 1e6, np.zeros((coords.shape[0], 1))), axis=1) coords_pt = torch.from_numpy(coords).int() feats_pt = torch.from_numpy(feats).float() labels_pt = torch.from_numpy(labels).long().to(device) input_pt = ME.SparseTensor( feats_pt, coords=coords_pt, quantization_mode=ME.SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE ).to(device) print(labels[labels == 1].shape) return input_pt, labels_pt, coords, feats, labels
def val(net, device, config, phase="val"): is_minknet = isinstance(net, ME.MinkowskiNetwork) data_loader = make_data_loader_custom( "val", config=config, ) net.eval() labels_val, preds_val = [], [] with torch.no_grad(): for batch in data_loader: coords, feats, labels = batch input = ME.SparseTensor(feats.float(), coords, device="cuda") logit = net(input) pred = torch.argmax(logit, 1) #print("val_labels", labels) #print("val_logit", logit) #print("val_pred", pred) #labels.append(labels.numpy()) labels_val.append(labels.cpu().numpy()) preds_val.append(pred.cpu().numpy()) torch.cuda.empty_cache() return metrics.accuracy_score(np.concatenate(labels_val), np.concatenate(preds_val))
def forward(self, batch): # Coords must be on CPU, features can be on GPU - see MinkowskiEngine documentation x = ME.SparseTensor(features=batch['features'], coordinates=batch['coords']) x = self.backbone(x) # x is (num_points, n_features) tensor assert x.shape[1] == self.feature_size, 'Backbone output tensor has: {} channels. Expected: {}'.format(x.shape[1], self.feature_size) x = self.pooling(x) if x.dim() == 3 and x.shape[2] == 1: # Reshape (batch_size, x = x.flatten(1) assert x.dim() == 2, 'Expected 2-dimensional tensor (batch_size,output_dim). Got {} dimensions.'.format(x.dim()) assert x.shape[1] == self.pooled_feature_size, 'Backbone output tensor has: {} channels. Expected: {}'.format(x.shape[1], self.pooled_feature_size) if self.dropout is not None: x = self.dropout(x) if self.linear is not None: x = self.linear(x) assert x.shape[1] == self.output_dim, 'Output tensor has: {} channels. Expected: {}'.format(x.shape[1], self.output_dim) # x is (batch_size, output_dim) tensor return {'embedding': x}
def training_step(self, batch, batch_idx): stensor = ME.SparseTensor(coordinates=batch["coordinates"], features=batch["features"]) # Must clear cache at regular interval if self.global_step % 10 == 0: torch.cuda.empty_cache() return self.criterion(self(stensor).F, batch["labels"].long())
def fcgf_feature_extraction(self, feats, coords): ''' Step 1: extract fast and accurate FCGF feature per point ''' sinput = ME.SparseTensor(feats, coords=coords).to(self.device) return self.fcgf_model(sinput).F
def forward(self, input): coords = input[:, 0:self.D + 1].cpu().int() features = input[:, self.D + 1:].float() x = ME.SparseTensor(features, coords=coords) x = self.input_layer(x) out = self.acnn(x) return out
def extract_feats(pcd, feature_type, voxel_size, model=None): xyz = np.asarray(pcd.points) _, sel = ME.utils.sparse_quantize(xyz, return_index=True, quantization_size=voxel_size) xyz = xyz[sel] pcd = make_o3d_pointcloud(xyz) if feature_type == 'FPFH': radius_normal = voxel_size * 2 pcd.estimate_normals( o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=30)) radius_feat = voxel_size * 5 feat = o3d.pipelines.registration.compute_fpfh_feature( pcd, o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feat, max_nn=100)) # (N, 33) return pcd, feat.data.T elif feature_type == 'FCGF': DEVICE = torch.device('cuda') coords = ME.utils.batched_coordinates( [torch.floor(torch.from_numpy(xyz) / voxel_size).int()]).to(DEVICE) feats = torch.ones(coords.size(0), 1).to(DEVICE) sinput = ME.SparseTensor(feats, coordinates=coords) # .to(DEVICE) # (N, 32) return pcd, model(sinput).F.detach().cpu().numpy() else: raise NotImplementedError( 'Unimplemented feature type {}'.format(feature_type))
def forward(self, x: ME.SparseTensor): # This implicitly applies ReLU on x (clamps negative values) temp = ME.SparseTensor(x.F.clamp(min=self.eps).pow(self.p), coordinates=x.C) temp = self.f(temp) # Apply ME.MinkowskiGlobalAvgPooling return temp.F.pow(1. / self.p) # Return (batch_size, n_features) tensor
def decomposition(): coords0, feats0 = to_sparse_coo(data_batch_0) coords1, feats1 = to_sparse_coo(data_batch_1) coords, feats = ME.utils.sparse_collate(coords=[coords0, coords1], feats=[feats0, feats1]) # sparse tensors A = ME.SparseTensor(coords=coords, feats=feats) conv = ME.MinkowskiConvolution(in_channels=1, out_channels=2, kernel_size=3, stride=2, dimension=2) B = conv(A) # Extract features and coordinates per batch index list_of_coords = B.decomposed_coordinates list_of_feats = B.decomposed_features list_of_coords, list_of_feats = B.decomposed_coordinates_and_features # To specify a batch index batch_index = 1 coords = B.coordinates_at(batch_index) feats = B.features_at(batch_index) # Empty list if given an invalid batch index batch_index = 3 print(B.coordinates_at(batch_index))
def sparse_tensor_arithmetics(): coords0, feats0 = to_sparse_coo(data_batch_0) coords0, feats0 = ME.utils.sparse_collate(coords=[coords0], feats=[feats0]) coords1, feats1 = to_sparse_coo(data_batch_1) coords1, feats1 = ME.utils.sparse_collate(coords=[coords1], feats=[feats1]) # sparse tensors A = ME.SparseTensor(coordinates=coords0, features=feats0) B = ME.SparseTensor(coordinates=coords1, features=feats1) # The following fails try: C = A + B except AssertionError: pass B = ME.SparseTensor( coordinates=coords1, features=feats1, coordinate_manager=A. coordinate_manager # must share the same coordinate manager ) C = A + B C = A - B C = A * B C = A / B # in place operations # Note that it requires the same coords_key (no need to feed coords) D = ME.SparseTensor( # coords=coords, not required features=feats0, coordinate_manager=A. coordinate_manager, # must share the same coordinate manager coordinate_map_key=A. coordinate_map_key # For inplace, must share the same coords key ) A += D A -= D A *= D A /= D # If you have two or more sparse tensors with the same coords_key, you can concatenate features E = ME.cat(A, D)
def generate_input_sparse_tensor(file_name, voxel_size=0.05): # Create a batch, this process is done in a data loader during training in parallel. batch = [load_file(file_name, voxel_size)] coordinates_, featrues_, pcds = list(zip(*batch)) coordinates, features = ME.utils.sparse_collate(coordinates_, featrues_) # Normalize features and create a sparse tensor return ME.SparseTensor(features - 0.5, coords=coordinates).to(device)
def inlier_prediction(self, inlier_feats, coords): ''' Step 4: predict inlier likelihood ''' sinput = ME.SparseTensor(inlier_feats, coords=coords).to(self.device) soutput = self.inlier_model(sinput) return soutput.F
def test_sum(self): coords, colors, pcd = load_file("1.ply") device = "cuda" D = 3 batch_size = 16 voxel_size = 0.02 channels = [3, 64, 128] dcoords = torch.from_numpy(np.floor(coords / voxel_size)).int() bcoords = batched_coordinates([dcoords for i in range(batch_size)]) in_feats = torch.rand(len(bcoords), 3).to(0) layer = MinkowskiStackSum( ME.MinkowskiConvolution( channels[0], channels[1], kernel_size=3, stride=1, dimension=3, ), nn.Sequential( ME.MinkowskiConvolution( channels[0], channels[1], kernel_size=3, stride=2, dimension=3, ), ME.MinkowskiStackSum( nn.Identity(), nn.Sequential( ME.MinkowskiConvolution( channels[1], channels[2], kernel_size=3, stride=2, dimension=3, ), ME.MinkowskiConvolutionTranspose( channels[2], channels[1], kernel_size=3, stride=1, dimension=3, ), ME.MinkowskiPoolingTranspose( kernel_size=2, stride=2, dimension=D ), ), ), ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D), ), ).cuda() for i in range(1000): torch.cuda.empty_cache() sinput = ME.SparseTensor(in_feats, coordinates=bcoords, device=device) layer(sinput)