def test_model_fn(batch, model, epoch): coords = batch['locs'].cuda() # (N, 1 + 3), long, cuda, dimension 0 for batch_idx voxel_coords = batch['voxel_locs'].cuda() # (M, 1 + 3), long, cuda p2v_map = batch['p2v_map'].cuda() # (N), int, cuda v2p_map = batch['v2p_map'].cuda() # (M, 1 + maxActive), int, cuda coords_float = batch['locs_float'].cuda() # (N, 3), float32, cuda feats = batch['feats'].cuda() # (N, C), float32, cuda batch_offsets = batch['offsets'].cuda() # (B + 1), int, cuda spatial_shape = batch['spatial_shape'] if cfg.use_coords: feats = torch.cat((feats, coords_float), 1) voxel_feats = pointgroup_ops.voxelization(feats, v2p_map, cfg.mode) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, voxel_coords.int(), spatial_shape, cfg.batch_size) ret = model(input_, p2v_map, coords_float, coords[:, 0].int(), batch_offsets, epoch) semantic_scores = ret['semantic_scores'] # (N, nClass) float32, cuda pt_offsets = ret['pt_offsets'] # (N, 3), float32, cuda if (epoch > cfg.prepare_epochs): scores, proposals_idx, proposals_offset = ret['proposal_scores'] ##### preds with torch.no_grad(): preds = {} preds['semantic'] = semantic_scores preds['pt_offsets'] = pt_offsets if (epoch > cfg.prepare_epochs): preds['score'] = scores preds['proposals'] = (proposals_idx, proposals_offset) return preds
def clusters_voxelization(self, clusters_idx, clusters_offset, feats, coords, fullscale, scale, mode): ''' :param clusters_idx: (SumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N, cpu :param clusters_offset: (nCluster + 1), int, cpu :param feats: (N, C), float, cuda :param coords: (N, 3), float, cuda :return: ''' c_idxs = clusters_idx[:, 1].cuda() clusters_feats = feats[c_idxs.long()] clusters_coords = coords[c_idxs.long()] clusters_coords_mean = pointgroup_ops.sec_mean(clusters_coords, clusters_offset.cuda()) # (nCluster, 3), float clusters_coords_mean = torch.index_select(clusters_coords_mean, 0, clusters_idx[:, 0].cuda().long()) # (sumNPoint, 3), float clusters_coords -= clusters_coords_mean clusters_coords_min = pointgroup_ops.sec_min(clusters_coords, clusters_offset.cuda()) # (nCluster, 3), float clusters_coords_max = pointgroup_ops.sec_max(clusters_coords, clusters_offset.cuda()) # (nCluster, 3), float clusters_scale = 1 / ((clusters_coords_max - clusters_coords_min) / fullscale).max(1)[0] - 0.01 # (nCluster), float clusters_scale = torch.clamp(clusters_scale, min=None, max=scale) min_xyz = clusters_coords_min * clusters_scale.unsqueeze(-1) # (nCluster, 3), float max_xyz = clusters_coords_max * clusters_scale.unsqueeze(-1) clusters_scale = torch.index_select(clusters_scale, 0, clusters_idx[:, 0].cuda().long()) clusters_coords = clusters_coords * clusters_scale.unsqueeze(-1) range = max_xyz - min_xyz offset = - min_xyz + torch.clamp(fullscale - range - 0.001, min=0) * torch.rand(3).cuda() + torch.clamp(fullscale - range + 0.001, max=0) * torch.rand(3).cuda() offset = torch.index_select(offset, 0, clusters_idx[:, 0].cuda().long()) clusters_coords += offset assert clusters_coords.shape.numel() == ((clusters_coords >= 0) * (clusters_coords < fullscale)).sum() clusters_coords = clusters_coords.long() clusters_coords = torch.cat([clusters_idx[:, 0].view(-1, 1).long(), clusters_coords.cpu()], 1) # (sumNPoint, 1 + 3) out_coords, inp_map, out_map = pointgroup_ops.voxelization_idx(clusters_coords, int(clusters_idx[-1, 0]) + 1, mode) # output_coords: M * (1 + 3) long # input_map: sumNPoint int # output_map: M * (maxActive + 1) int out_feats = pointgroup_ops.voxelization(clusters_feats, out_map.cuda(), mode) # (M, C), float, cuda spatial_shape = [fullscale] * 3 voxelization_feats = spconv.SparseConvTensor(out_feats, out_coords.int().cuda(), spatial_shape, int(clusters_idx[-1, 0]) + 1) return voxelization_feats, inp_map
def point_voxelization(self, data_dict): '''Transfering batch point clouds to a sparse tensor Args: data_dict: dict Returns: input:SparseTensor ''' feats = data_dict['seg_features'].cuda() v2p_map = data_dict['v2p_map'].cuda() voxel_locs = data_dict['voxel_locs'].int().cuda() spatial_shape = data_dict['spatial_shape'] voxel_feats = pointgroup_ops.voxelization( feats, v2p_map, self.config['Segmentation']['mode']) # (M, C), float, cuda input = spconv.SparseConvTensor(voxel_feats, voxel_locs, spatial_shape, self.config['TRAIN']['batch_size']) return input
def model_fn(batch, model, epoch): #print ('model fn') ##### prepare input and forward # batch {'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, # 'locs_float': locs_float, 'feats': feats, 'labels': labels, 'instance_labels': instance_labels, # 'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, # 'id': tbl, 'offsets': batch_offsets, 'spatial_shape': spatial_shape} coords = batch['locs'].cuda( ) # (N, 1 + 3), long, cuda, dimension 0 for batch_idx voxel_coords = batch['voxel_locs'].cuda() # (M, 1 + 3), long, cuda p2v_map = batch['p2v_map'].cuda() # (N), int, cuda v2p_map = batch['v2p_map'].cuda() # (M, 1 + maxActive), int, cuda coords_float = batch['locs_float'].cuda() # (N, 3), float32, cuda feats = batch['feats'].cuda() # (N, C), float32, cuda labels = batch['labels'].cuda() # (N), long, cuda '''instance_labels = batch['instance_labels'].cuda() # (N), long, cuda, 0~total_nInst, -100 instance_info = batch['instance_info'].cuda() # (N, 9), float32, cuda, (meanxyz, minxyz, maxxyz) instance_pointnum = batch['instance_pointnum'].cuda() # (total_nInst), int, cuda''' batch_offsets = batch['offsets'].cuda() # (B + 1), int, cuda spatial_shape = batch['spatial_shape'] if cfg.use_coords: feats = torch.cat((feats, coords_float), 1) voxel_feats = pointgroup_ops.voxelization( feats, v2p_map, cfg.mode) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, voxel_coords.int(), spatial_shape, cfg.batch_size) ret = model(input_, p2v_map, coords_float, coords[:, 0].int(), batch_offsets, epoch) semantic_scores = ret['semantic_scores'] # (N, nClass) float32, cuda '''pt_offsets = ret['pt_offsets'] # (N, 3), float32, cuda if(epoch > cfg.prepare_epochs): scores, proposals_idx, proposals_offset = ret['proposal_scores'] # scores: (nProposal, 1) float, cuda # proposals_idx: (sumNPoint, 2), int, cpu, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # proposals_offset: (nProposal + 1), int, cpu''' loss_inp = {} loss_inp['semantic_scores'] = (semantic_scores, labels) '''loss_inp['pt_offsets'] = (pt_offsets, coords_float, instance_info, instance_labels) if(epoch > cfg.prepare_epochs): loss_inp['proposal_scores'] = (scores, proposals_idx, proposals_offset, instance_pointnum)''' loss, loss_out, infos = loss_fn(loss_inp, epoch) ##### accuracy / visual_dict / meter_dict with torch.no_grad(): preds = {} preds['semantic'] = semantic_scores '''preds['pt_offsets'] = pt_offsets if(epoch > cfg.prepare_epochs): preds['score'] = scores preds['proposals'] = (proposals_idx, proposals_offset)''' visual_dict = {} visual_dict['loss'] = loss for k, v in loss_out.items(): visual_dict[k] = v[0] meter_dict = {} meter_dict['loss'] = (loss.item(), coords.shape[0]) for k, v in loss_out.items(): meter_dict[k] = (float(v[0]), v[1]) #print (meter_dict.keys()) return loss, preds, visual_dict, meter_dict
def model_fn(batch, model, epoch): #print ('model fn') ##### prepare input and forward # batch {'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, # 'locs_float': locs_float, 'feats': feats, 'labels': labels, 'instance_labels': instance_labels, # 'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, # 'id': tbl, 'offsets': batch_offsets, 'spatial_shape': spatial_shape} coords = batch['locs'].cuda( ) # (N, 1 + 3), long, cuda, dimension 0 for batch_idx voxel_coords = batch['voxel_locs'].cuda() # (M, 1 + 3), long, cuda p2v_map = batch['p2v_map'].cuda() # (N), int, cuda v2p_map = batch['v2p_map'].cuda() # (M, 1 + maxActive), int, cuda coords_float = batch['locs_float'].cuda() # (N, 3), float32, cuda feats = batch['feats'].cuda() # (N, C), float32, cuda labels = batch['labels'].cuda() # (N), long, cuda groups = batch['groups'] group2points = batch['group2points'] group_fulls = batch['group_fulls'] group2point_fulls = batch['group2point_fulls'] #for i in range(4): # print (len(groups[i][0])) classes = [] poss = [] negs = [] for i in range(20): classes.append([]) for g in range(len(groups)): group = groups[g] for i in range(20): for s in range(len(group[i])): classes[i].append((i, g, group[i][s])) ignore = [] mini = 10 #min(min(map(len, classes)),30) for i in range(20): random.shuffle(classes[i]) if len(classes[i]) == 0: ignore.append(i) continue if len(classes[i]) >= mini: classes[i] = classes[i][:mini] else: while (len(classes[i]) < 10): classes[i].append(random.choice(classes[i])) batch_offsets = batch['offsets'].cuda() # (B + 1), int, cuda spatial_shape = batch['spatial_shape'] if cfg.use_coords: feats = torch.cat((feats, coords_float), 1) voxel_feats = pointgroup_ops.voxelization( feats, v2p_map, cfg.mode) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, voxel_coords.int(), spatial_shape, cfg.batch_size) ret, output_feats = model(input_, p2v_map, coords_float, coords[:, 0].int(), batch_offsets, epoch) tmpfeat = torch.Tensor(20, mini, model.m).cuda() label = torch.zeros(20 * mini).long().cuda() for i in range(20): if i in ignore: tmpfeat[i, :, :] = 0 label[i * mini:i * mini + mini] = -100 continue for j in range(mini): sample = classes[i][j] c0 = sample[0] b0 = sample[1] idx0 = sample[2] idx_off = torch.tensor(np.asarray( group2points[b0][idx0])) + batch['offsets'][b0] feat = output_feats[idx_off] feat = torch.mean(feat, 0) tmpfeat[i, j, :] = feat label[i * mini + j] = i if model.start == 1: model.start = 0 model.feat = torch.mean(tmpfeat.detach(), 1) else: model.feat = 0.9 * model.feat + 0.1 * torch.mean( tmpfeat.detach(), 1) #model.feat=nn.functional.normalize(model.feat,1) #tmpfeat=nn.functional.normalize(tmpfeat,2) tmpfeat = torch.reshape(tmpfeat, (20 * mini, model.m)) product = torch.matmul(tmpfeat, torch.transpose(model.feat, 0, 1)) / 0.07 semantic_scores_pred = ret[ 'semantic_scores'] # (N, nClass) float32, cuda semantic_scores_crfs = torch.zeros( (semantic_scores_pred.shape[0], 20)).cuda() result_crfs = torch.zeros((semantic_scores_pred.shape[0], )).cuda() #print ('semantic_scores_crfs',semantic_scores_crfs.shape) for batch_idx in range(len(groups)): group = groups[batch_idx] start = batch['offsets'][batch_idx] end = batch['offsets'][batch_idx + 1] #CRF Q = torch.nn.functional.softmax( semantic_scores_pred[start:end, :], 1) #torch.exp(semantic_scores_pred) #print (start,end,Q.shape) cnt = 0 group_full = group_fulls[batch_idx] pair = torch.zeros((len(group_full), 6 + 32)).cuda() U = torch.zeros((len(group_full), 20)).cuda() for i in group_full: idxs = torch.from_numpy( np.asarray(group2point_fulls[batch_idx][i])) #.cuda() pair[cnt, :6] = torch.mean(feats[idxs], 0) * 1 #pair[cnt,3:6]=torch.mean(coords_float[idxs],0)*0 pair[cnt, 6:38] = torch.mean(output_feats[idxs], 0) * 0 #print (Q[idxs],Q[idxs].shape) U[cnt, :] = torch.mean(Q[idxs], 0) #print (torch.max(idxs)) assert (torch.max(idxs) < Q.shape[0]) #logits[cnt,]=torch.mean(semantic_scores_pred[idxs],0) cnt += 1 #U=U/torch.abs(torch.sum(U,1)) #print ('U', U,U.shape,torch.unique(U)) #Usum=torch.unsqueeze(torch.abs(torch.sum(U,1)),1).repeat(1,20) #print (torch.abs(torch.sum(U,1))) #U=U/Usum #print (U,U.shape,torch.unique(U)) #assert Q_ori = torch.log((U + 1e-5)) #print ('Q_ori',name,torch.unique(Q_ori)) pair1 = torch.unsqueeze(pair, 1).repeat(1, pair.shape[0], 1) pair2 = torch.transpose(pair1, 0, 1) #print (U1.shape,U2.shape) diff = torch.sum((pair1 - pair2).pow(2), 2) #print ('torch.exp(-0.5*diff)',torch.exp(-0.5*diff),torch.unique(torch.exp(-0.5*diff)),torch.exp(-0.5*diff).shape) k = torch.exp(-0.5 * diff) #w*exp(-0.5 * |f_i - f_j|^2) Q = U for it in range(1): #print ('iter',it) assert k.shape[0] > 0 kQ = torch.mm(k, Q) / k.shape[0] * 1 #kQ=torch.mm(kQ,miu) #Qsum=torch.expand_dims(torch.sum(kQ,1),1,20)-kQ #print ('Q',torch.unique(Q)) #print ('kQ',kQ,torch.unique(kQ),kQ.shape) #print ('expkQ',torch.exp(kQ)) #kQ[:]=0 #print ('Q_ori',Q_ori,torch.unique(Q_ori),Q_ori.shape) newQ = Q_ori + kQ #target is: Q_ori*torch.exp(kQ) #newQ=torch.nn.functional.softmax(newQ,1) #print ('nq', torch.unique(newQ)) Q = newQ #.clone() #print ('newQ',newQ,torch.unique(newQ),newQ.shape,torch.unique(torch.sum(newQ,1))) #Q_result=torch.argmax(Q,1) semantic_scores_crf = torch.zeros((end - start, 20)).cuda() semantic_scores_crf[:] = semantic_scores_pred[start: end, :] #.clone() #result_crf=torch.zeros((semantic_scores_pred.shape[0],)) #print (group_full,len(group_full))' #print (Q.shape) for i in range(len(group_full)): #print (i) g = group_full[i] #print (g) #print (group2point_fulls[batch_idx]) #print (group2point_fulls[batch_idx][g]) #print (np.asarray(group2point_fulls[batch_idx][g])) idxs = torch.from_numpy( np.asarray(group2point_fulls[batch_idx][g])).cuda() #print (idxs)] #tmp=torch.unsqueeze(Q[i],0) #print ('Q',tmp,torch.unique(tmp),tmp.shape) #print (i,newQ[i]) #for j in idxs: semantic_scores_crf[idxs, :] = newQ[i] #result_crf[idxs]=torch.argmax(Q[i],1) #print (semantic_scores_crf.shape, torch.unique(semantic_scores_crf)) semantic_scores_crfs[start:end, :] = semantic_scores_crf #print (start,end) #print ('semantic_scros_crf',semantic_scores_crfs.shape,torch.unique(semantic_scores_crfs)) loss_inp = {} loss_inp['semantic_scores'] = (product, label) loss_inp['semantic_scores_crf'] = (semantic_scores_crfs, labels) loss_inp['semantic_scores_pred'] = (semantic_scores_pred, labels) #print (semantic_scores_crfs, labels, semantic_scores_pred, labels) loss, loss_relation, loss_pred, loss_crf, loss_out, infos = loss_fn( loss_inp, epoch) ##### accuracy / visual_dict / meter_dict with torch.no_grad(): preds = {} #preds['semantic'] = semantic_scores visual_dict = {} visual_dict['loss'] = loss for k, v in loss_out.items(): visual_dict[k] = v[0] meter_dict = {} meter_dict['loss'] = (loss.item(), coords.shape[0]) for k, v in loss_out.items(): meter_dict[k] = (float(v[0]), v[1]) return loss, loss_relation, loss_pred, loss_crf, preds, visual_dict, meter_dict
def forward(self, input, input_map, coords, rgb, ori_coords, batch_idxs, batch_offsets, epoch): ''' :param input_map: (N), int, cuda :param coords: (N, 3), float, cuda :param batch_idxs: (N), int, cuda :param batch_offsets: (B + 1), int, cuda ''' ret = {} batch_idxs = batch_idxs.squeeze() semantic_scores = [] point_offset_preds = [] voxel_occupancy_preds = [] voxel_feats = pointgroup_ops.voxelization( input['pt_feats'], input['v2p_map'], input['mode']) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, input['voxel_coords'], input['spatial_shape'], input['batch_size']) output = self.input_conv(input_) output = self.unet(output) output = self.output_layer(output) output_feats = output.features[input_map.long()] output_feats = output_feats.squeeze(dim=0) ### point prediction #### point semantic label prediction semantic_scores.append( self.point_semantic(output_feats)) # (N, nClass), float ### only used to evaluate based on ground truth # semantic_scores.append(input['point_semantic_scores'][0]) # (N, nClass), float ### ground truth for each category # CATE_NUM = 0 # semantic_output = self.point_semantic(output_feats) # if (input['point_semantic_scores'][0].max(dim=1)[1] == CATE_NUM).sum() > 0: # semantic_output[input['point_semantic_scores'][0].max(dim=1)[1] == CATE_NUM] = \ # input['point_semantic_scores'][0][input['point_semantic_scores'][0].max(dim=1)[1] == CATE_NUM].float() # semantic_output[semantic_output.max(dim=1)[1] == CATE_NUM] = \ # input['point_semantic_scores'][0][semantic_output.max(dim=1)[1] == CATE_NUM].float() # semantic_scores.append(semantic_output) point_semantic_preds = semantic_scores[0].max(1)[1] #### point offset prediction point_offset_pred = self.point_offset(output_feats) point_offset_preds.append(point_offset_pred) # (N, 3), float32 # only used to evaluate based on ground truth # point_offset_preds.append(input['point_offset_preds']) # (N, 3), float32 voxel_occupancy_preds.append(self.point_occupancy(output.features)) point_occupancy_pred = voxel_occupancy_preds[0][ input_map.long()].squeeze(dim=1) if (epoch > self.prepare_epochs): #### get prooposal clusters object_idxs = torch.nonzero(point_semantic_preds > 1).view(-1) batch_idxs_ = batch_idxs[object_idxs] batch_offsets_ = utils.get_batch_offsets(batch_idxs_, input['batch_size']) coords_ = coords[object_idxs] pt_offsets_ = point_offset_preds[0][object_idxs] point_occupancy_pred_ = point_occupancy_pred[object_idxs] semantic_preds_cpu = point_semantic_preds[object_idxs].int().cpu() # idx_occupancy, start_len_occupancy = pointgroup_ops.ballquery_batch_p( # coords_ + pt_offsets_, batch_idxs_, # batch_offsets_, self.cluster_radius, self.cluster_shift_meanActive # ) # proposals_idx_occupancy, proposals_offset_occupancy = pointgroup_ops.bfs_occupancy_cluster( # semantic_preds_cpu, point_occupancy_pred_.cpu(), idx_occupancy.cpu(), # start_len_occupancy.cpu(), self.cluster_npoint_thre, self.occupancy_cluster['occupancy_threshold_shift'] # ) # proposals_idx_occupancy[:, 1] = object_idxs[proposals_idx_occupancy[:, 1].long()].int() # # proposals_idx_shift: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # # proposals_offset_shift: (nProposal + 1), int # # idx, start_len = pointgroup_ops.ballquery_batch_p( # coords_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_meanActive # ) # proposals_idx, proposals_offset = pointgroup_ops.bfs_occupancy_cluster( # semantic_preds_cpu, point_occupancy_pred_.cpu(), idx.cpu(), # start_len.cpu(), self.cluster_npoint_thre, self.occupancy_cluster['occupancy_threshold'] # ) # proposals_idx[:, 1] = object_idxs[proposals_idx[:, 1].long()].int() idx, start_len = pointgroup_ops.ballquery_batch_p( coords_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_meanActive) proposals_idx, proposals_offset = pointgroup_ops.bfs_cluster( semantic_preds_cpu, idx.cpu(), start_len.cpu(), self.cluster_npoint_thre) proposals_idx[:, 1] = object_idxs[proposals_idx[:, 1].long()].int() idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p( coords_ + pt_offsets_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_shift_meanActive) proposals_idx_shift, proposals_offset_shift = pointgroup_ops.bfs_cluster( semantic_preds_cpu, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre) proposals_idx_shift[:, 1] = object_idxs[ proposals_idx_shift[:, 1].long()].int() # proposals_idx_shift: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # proposals_offset_shift: (nProposal + 1), int # proposals_idx_occupancy[:, 0] += (proposals_offset.size(0) - 1) # proposals_offset_occupancy += proposals_offset[-1] # proposals_idx = torch.cat((proposals_idx, proposals_idx_occupancy), dim=0) # proposals_offset = torch.cat((proposals_offset, proposals_offset_occupancy[1:])) # proposals_idx_filtered = [] # proposals_offset_filtered = [0] # for proposal in proposals_idx_shift[:, 0].unique(): # proposal_index = proposals_idx_shift[:, 0] == proposal # proposal_idxs = proposals_idx_shift[proposal_index, 1].long() # proposals_indexs = proposals_idx_shift[proposal_index, :] # # proposal_occupancy_mean = point_occupancy_pred[proposal_idxs].mean() # # valid_point_index = torch.ones_like(proposal_idxs).byte() # valid_point_index[point_occupancy_pred[proposal_idxs] < proposal_occupancy_mean * # (1 - self.occupancy_cluster['occupancy_filter_threshold'])] = 0 # valid_point_index[point_occupancy_pred[proposal_idxs] > proposal_occupancy_mean * # (1 + self.occupancy_cluster['occupancy_filter_threshold'])] = 0 # proposal_idx_filtered = proposals_indexs[valid_point_index, :] # # proposals_idx_filtered.append(proposal_idx_filtered) # proposals_offset_filtered.append(proposals_offset_filtered[-1] + proposal_idx_filtered.shape[0]) # # proposals_idx_filtered = torch.cat(proposals_idx_filtered, dim=0) # proposals_offset_filtered = torch.tensor(proposals_offset_filtered).int() # proposals_idx_shift = proposals_idx_filtered # proposals_offset_shift = proposals_offset_filtered proposals_idx_shift[:, 0] += (proposals_offset.size(0) - 1) proposals_offset_shift += proposals_offset[-1] proposals_idx = torch.cat((proposals_idx, proposals_idx_shift), dim=0) proposals_offset = torch.cat( (proposals_offset, proposals_offset_shift[1:])) #### proposals voxelization again input_feats, inp_map = self.clusters_voxelization( proposals_idx, proposals_offset, output_feats, coords, self.score_fullscale, self.score_scale, self.mode) #### score score = self.score_unet(input_feats) score = self.score_outputlayer(score) score_feats = score.features[inp_map.long()] # (sumNPoint, C) score_feats = pointgroup_ops.roipool( score_feats, proposals_offset.cuda()) # (nProposal, C) scores = self.score_linear(score_feats) # (nProposal, 1) # proposals_idx = proposals_idx_occupancy # proposals_offset = proposals_offset_occupancy # scores = torch.ones(proposals_offset.shape[0] - 1, 1).to(point_offset_preds[0].device) ret['proposal_scores'] = (scores, proposals_idx, proposals_offset) ret['point_semantic_scores'] = semantic_scores ret['point_offset_preds'] = point_offset_preds ret['point_features'] = output_feats ret['voxel_occupancy_preds'] = voxel_occupancy_preds return ret
def forward(self, input, input_map, coords, rgb, ori_coords, batch_idxs, batch_offsets, epoch): ''' :param input_map: (N), int, cuda :param coords: (N, 3), float, cuda :param batch_idxs: (N), int, cuda :param batch_offsets: (B + 1), int, cuda ''' ret = {} batch_idxs = batch_idxs.squeeze() semantic_scores = [] point_offset_preds = [] voxel_feats = pointgroup_ops.voxelization( input['pt_feats'], input['v2p_map'], input['mode']) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, input['voxel_coords'], input['spatial_shape'], input['batch_size']) output = self.input_conv(input_) output = self.unet(output) output = self.output_layer(output) output_feats = output.features[input_map.long()] output_feats = output_feats.squeeze(dim=0) ### point prediction #### point semantic label prediction semantic_scores.append( self.point_semantic(output_feats)) # (N, nClass), float ### only used to evaluate based on ground truth # semantic_scores.append(input['point_semantic_scores'][0]) # (N, nClass), float ### ground truth for each category # CATE_NUM = 0 # semantic_output = self.point_semantic(output_feats) # if (input['point_semantic_scores'][0].max(dim=1)[1] == CATE_NUM).sum() > 0: # semantic_output[input['point_semantic_scores'][0].max(dim=1)[1] == CATE_NUM] = \ # input['point_semantic_scores'][0][input['point_semantic_scores'][0].max(dim=1)[1] == CATE_NUM].float() # semantic_output[semantic_output.max(dim=1)[1] == CATE_NUM] = \ # input['point_semantic_scores'][0][semantic_output.max(dim=1)[1] == CATE_NUM].float() # semantic_scores.append(semantic_output) point_semantic_preds = semantic_scores[0].max(1)[1] #### point offset prediction point_offset_pred = self.point_offset(output_feats) if self.instance_triplet_loss['activate']: point_offset_pred = point_offset_pred - input['pt_feats'][:, 3:] point_offset_preds.append(point_offset_pred) # (N, 3), float32 # only used to evaluate based on ground truth # point_offset_preds.append(input['point_offset_preds']) # (N, 3), float32 if self.voxel_center_prediction['activate']: voxel_center_preds = self.voxel_center_pred(output.features) voxel_center_offset_preds = self.voxel_center_offset( output.features) voxel_center_semantic_preds = self.voxel_center_semantic( output.features) ret['voxel_center_preds'] = voxel_center_preds ret['voxel_center_offset_preds'] = voxel_center_offset_preds ret['voxel_center_semantic_preds'] = voxel_center_semantic_preds if self.point_xyz_reconstruction_loss['activate']: point_reconstructed_coords = self.point_reconstruction_coords( output_feats) ret['point_reconstructed_coords'] = point_reconstructed_coords if self.instance_classifier['activate']: instance_id_preds = self.point_instance_classifier(output_feats) ret['instance_id_preds'] = instance_id_preds if (epoch > self.prepare_epochs): #### get prooposal clusters object_idxs = torch.nonzero(point_semantic_preds > 1).view(-1) batch_idxs_ = batch_idxs[object_idxs] batch_offsets_ = utils.get_batch_offsets(batch_idxs_, input['batch_size']) coords_ = coords[object_idxs] pt_offsets_ = point_offset_preds[0][object_idxs] semantic_preds_cpu = point_semantic_preds[object_idxs].int().cpu() idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p( coords_ + pt_offsets_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_shift_meanActive) # idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p( # coords_ + pt_offsets_ + (torch.rand(coords_.shape) * 1e-2).cuda(), batch_idxs_, # batch_offsets_, 0.001, self.cluster_shift_meanActive # ) proposals_idx_shift, proposals_offset_shift = pointgroup_ops.bfs_cluster( semantic_preds_cpu, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre) proposals_idx_shift[:, 1] = object_idxs[ proposals_idx_shift[:, 1].long()].int() # proposals_idx_shift: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # proposals_offset_shift: (nProposal + 1), int idx, start_len = pointgroup_ops.ballquery_batch_p( coords_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_meanActive) proposals_idx, proposals_offset = pointgroup_ops.bfs_cluster( semantic_preds_cpu, idx.cpu(), start_len.cpu(), self.cluster_npoint_thre) proposals_idx[:, 1] = object_idxs[proposals_idx[:, 1].long()].int() # proposals_idx: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # proposals_offset: (nProposal + 1), int proposals_idx_shift[:, 0] += (proposals_offset.size(0) - 1) proposals_offset_shift += proposals_offset[-1] proposals_idx = torch.cat((proposals_idx, proposals_idx_shift), dim=0) proposals_offset = torch.cat( (proposals_offset, proposals_offset_shift[1:])) #### proposals voxelization again input_feats, inp_map = self.clusters_voxelization( proposals_idx, proposals_offset, output_feats, coords, self.score_fullscale, self.score_scale, self.mode) #### score score = self.score_unet(input_feats) score = self.score_outputlayer(score) score_feats = score.features[inp_map.long()] # (sumNPoint, C) score_feats = pointgroup_ops.roipool( score_feats, proposals_offset.cuda()) # (nProposal, C) scores = self.score_linear(score_feats) # (nProposal, 1) ret['proposal_scores'] = (scores, proposals_idx, proposals_offset) ret['point_semantic_scores'] = semantic_scores ret['point_offset_preds'] = point_offset_preds if self.instance_triplet_loss['activate']: ret['point_offset_feats'] = output_feats ret['point_features'] = output_feats return ret
def model_fn(batch, model, epoch): #print ('model fn') ##### prepare input and forward # batch {'locs': locs, 'voxel_locs': voxel_locs, 'p2v_map': p2v_map, 'v2p_map': v2p_map, # 'locs_float': locs_float, 'feats': feats, 'labels': labels, 'instance_labels': instance_labels, # 'instance_info': instance_infos, 'instance_pointnum': instance_pointnum, # 'id': tbl, 'offsets': batch_offsets, 'spatial_shape': spatial_shape} coords = batch['locs'].cuda( ) # (N, 1 + 3), long, cuda, dimension 0 for batch_idx voxel_coords = batch['voxel_locs'].cuda() # (M, 1 + 3), long, cuda p2v_map = batch['p2v_map'].cuda() # (N), int, cuda v2p_map = batch['v2p_map'].cuda() # (M, 1 + maxActive), int, cuda coords_float = batch['locs_float'].cuda() # (N, 3), float32, cuda feats = batch['feats'].cuda() # (N, C), float32, cuda labels = batch['labels'].cuda() # (N), long, cuda groups = batch['groups'] group2points = batch['group2points'] #for i in range(4): # print (len(groups[i][0])) classes = [] poss = [] negs = [] for i in range(20): classes.append([]) for g in range(len(groups)): group = groups[g] for i in range(20): for s in range(len(group[i])): classes[i].append((i, g, group[i][s])) ignore = [] mini = 10 #min(min(map(len, classes)),30) for i in range(20): random.shuffle(classes[i]) if len(classes[i]) == 0: ignore.append(i) continue if len(classes[i]) >= mini: classes[i] = classes[i][:mini] else: while (len(classes[i]) < 10): classes[i].append(random.choice(classes[i])) '''for i in range(20): if len(classes[i])==0: continue for times in range(10): poss.append((random.choice(classes[i]),random.choice(classes[i]))) for j in range(20): if j==i or len(classes[j])==0: continue negs.append((random.choice(classes[i]),random.choice(classes[j]))) while(len(negs)<20): j=random.randint(0,19) if j==i or len(classes[j])==0: continue negs.append((random.choice(classes[i]),random.choice(classes[j]))) posidxs=[] for pos in poss: c0 = pos[0][0] b0 = pos[0][1] idx0 = pos[0][2] #print ('fffff',group2points[b0].keys()) idx_off0 = torch.tensor(np.asarray(group2points[b0][idx0])) + batch['offsets'][b0] #print ('sssssss',group2points[b0]) c1 = pos[1][0] b1 = pos[1][1] idx1 = pos[1][2] idx_off1 = torch.tensor(np.asarray(group2points[b1][idx1])) + batch['offsets'][b1] posidxs.append((idx_off0,idx_off1)) negidxs=[] for neg in negs: c0 = neg[0][0] b0 = neg[0][1] idx0 = neg[0][2] idx_off0 = torch.tensor(np.asarray(group2points[b0][idx0])) + batch['offsets'][b0] #print (idx_off0) c1 = neg[1][0] b1 = neg[1][1] idx1 = neg[1][2] idx_off1 = torch.tensor(np.asarray(group2points[b1][idx1])) + batch['offsets'][b1] negidxs.append((idx_off0,idx_off1))''' batch_offsets = batch['offsets'].cuda() # (B + 1), int, cuda spatial_shape = batch['spatial_shape'] if cfg.use_coords: feats = torch.cat((feats, coords_float), 1) voxel_feats = pointgroup_ops.voxelization( feats, v2p_map, cfg.mode) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, voxel_coords.int(), spatial_shape, cfg.batch_size) ret, output_feats = model(input_, p2v_map, coords_float, coords[:, 0].int(), batch_offsets, epoch) #print (torch.unique(voxel_feats),output_feats) tmpfeat = torch.Tensor(20, mini, model.m).cuda() label = torch.zeros(20 * mini).long().cuda() for i in range(20): if i in ignore: tmpfeat[i, :, :] = 0 label[i * mini:i * mini + mini] = -100 continue for j in range(mini): sample = classes[i][j] c0 = sample[0] b0 = sample[1] idx0 = sample[2] #print ('fffff',group2points[b0].keys()) idx_off = torch.tensor(np.asarray( group2points[b0][idx0])) + batch['offsets'][b0] feat = output_feats[idx_off] feat = torch.mean(feat, 0) tmpfeat[i, j, :] = feat label[i * mini + j] = i if model.start == 1: model.start = 0 model.feat = torch.mean(tmpfeat.detach(), 1) else: model.feat = 0.9 * model.feat + 0.1 * torch.mean( tmpfeat.detach(), 1) #model.feat=nn.functional.normalize(model.feat,1) #tmpfeat=nn.functional.normalize(tmpfeat,2) tmpfeat = torch.reshape(tmpfeat, (20 * mini, model.m)) product = torch.matmul(tmpfeat, torch.transpose(model.feat, 0, 1)) / 0.07 loss_inp = {} loss_inp['semantic_scores'] = (product, label) '''posfeat=torch.zeros((10,128*2)) for j in range(len(posidxs)): posidx=posidxs[j] feat0=output_feats[posidx[0]] feat0=torch.mean(feat0,0) feat1=output_feats[posidx[1]] feat1=torch.mean(feat0,0) feat=torch.cat((feat0,feat1),1) posfeat[j][:128]=feat0 posfeat[j][128:]=feat1 negfeat=torch.zeros((20,128*2)) for j in range(len(negidxs)): negidx=negidxs[j] feat0=output_feats[negidx[0]] feat0=torch.mean(feat0,0) feat1=output_feats[negidx[1]] feat1=torch.mean(feat0,0) feat=torch.cat((feat0,feat1),1) negfeat[j][:128]=feat0 negfeat[j][128:]=feat1''' '''semantic_scores = ret['semantic_scores'] # (N, nClass) float32, cuda print ('s',semantic_scores.shape) #print ('groups',groups) loss_inp = {} loss_inp['semantic_scores'] = (semantic_scores, labels) print ('l',labels.shape,'\n')''' loss, loss_out, infos = loss_fn(loss_inp, epoch) ##### accuracy / visual_dict / meter_dict with torch.no_grad(): preds = {} #preds['semantic'] = semantic_scores visual_dict = {} visual_dict['loss'] = loss for k, v in loss_out.items(): visual_dict[k] = v[0] meter_dict = {} meter_dict['loss'] = (loss.item(), coords.shape[0]) for k, v in loss_out.items(): meter_dict[k] = (float(v[0]), v[1]) return loss, preds, visual_dict, meter_dict
def test_model_fn(batch, model, epoch): #print ('test model fn') feats = torch.zeros((20, 16)) coords = batch['locs'].cuda( ) # (N, 1 + 3), long, cuda, dimension 0 for batch_idx voxel_coords = batch['voxel_locs'].cuda() # (M, 1 + 3), long, cuda p2v_map = batch['p2v_map'].cuda() # (N), int, cuda v2p_map = batch['v2p_map'].cuda() # (M, 1 + maxActive), int, cuda coords_float = batch['locs_float'].cuda() # (N, 3), float32, cuda feats = batch['feats'].cuda() # (N, C), float32, cuda batch_offsets = batch['offsets'].cuda() # (B + 1), int, cuda spatial_shape = batch['spatial_shape'] labels = batch['labels'] groups = batch['groups'] group2points = batch['group2points'] if cfg.use_coords: feats = torch.cat((feats, coords_float), 1) voxel_feats = pointgroup_ops.voxelization( feats, v2p_map, cfg.mode) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, voxel_coords.int(), spatial_shape, cfg.batch_size) ret, output_feats = model(input_, p2v_map, coords_float, coords[:, 0].int(), batch_offsets, epoch) products = torch.zeros((output_feats.shape[0], 20)) result = torch.zeros((output_feats.shape[0])) result_feat = torch.zeros((output_feats.shape[0], 32)) #for group in groups: #print ('len',len(groups),groups) feat_voxel = [] gt_voxel = [] rgb_voxel = [] group = groups[0] for i in range(20): for s in range(len(group[i])): idxs = group2points[0][group[i][s]] feats = output_feats[idxs] #rgbs=feats[idxs] feat = torch.mean(feats, 0) #rgb=torch.mean(rgbs,0) #feat_voxel.append(feat.detach().cpu().numpy()) #gt_voxel.append(i) #rgb_voxel.append(rgb.detach().cpu().numpy()) #stdfeat=torch.Tensor(torch.load('exp/scannetv2/pointgroup/pointgroup_run1_scannet/pointgroup_run1_scannet-000000074_feat.pth')).cuda() product = torch.matmul(feat, torch.transpose(model.feat, 0, 1)) #print (i,int(torch.argmax(product)),product) #result[idxs]=int(torch.argmax(product)) for idx in idxs: result[idx] = int(torch.argmax(product)) result_feat[idx, :] = feat products[idx, :] = product #print (int(torch.argmax(product))) #print ('idxs',len(idxs),idxs) #print ('unique',torch.unique(result)) #semantic_scores=torch.argmax(product,1) '''feats=np.zeros((20,16)) labels=np.zeros((20)) for i in range(feat.shape[0]): if label[i]==-100: continue feats[label[i],:]+=feat[i,:].cpu().numpy() labels[label[i]]+=1 for i in range(20): feats[i,:]/=labels[i]''' #print (feat.shape,label.shape) #semantic_scores = ret['semantic_scores'] # (N, nClass) float32, cuda '''pt_offsets = ret['pt_offsets'] # (N, 3), float32, cuda if (epoch > cfg.prepare_epochs): scores, proposals_idx, proposals_offset = ret['proposal_scores']''' '''product=torch.matmul(feat,torch.transpose(model.feat,0,1)) print ('product',product.shape) result=torch.argmax(product,1) print ('result',result.shape)''' #loss_inp = {} #loss_inp['semantic_scores'] = (product, label) ##### preds with torch.no_grad(): preds = {} preds['semantic'] = result preds['products'] = products print('feat') preds['feat'] = result_feat #np.asarray(feat_voxel) #preds['gt'] = np.asarray(gt_voxel) #preds['rgb'] = np.asarray(rgb_voxel) '''preds['pt_offsets'] = pt_offsets if (epoch > cfg.prepare_epochs): preds['score'] = scores preds['proposals'] = (proposals_idx, proposals_offset)''' return preds
def forward(self, data_dict, epoch): ''' :param input_map: (N), int, cuda :param coords: (N, 3), float, cuda :param batch_idxs: (N), int, cuda :param batch_offsets: (B + 1), int, cuda ''' #unpack input data coords = data_dict[ 'locs'] # (N, 1 + 3), long, cuda, dimension 0 for batch_idx voxel_coords = data_dict['voxel_locs'] # (M, 1 + 3), long, cuda p2v_map = data_dict['p2v_map'] # (N), int, cuda v2p_map = data_dict['v2p_map'] # (M, 1 + maxActive), int, cuda coords_float = data_dict['locs_float'] # (N, 3), float32, cuda feats = data_dict['feats'] # (N, C), float32, cuda spatial_shape = data_dict['spatial_shape'] batch_idxs = coords[:, 0].int() if self.cfg.use_coords: feats = torch.cat((coords_float, feats), 1) voxel_feats = pointgroup_ops.voxelization( feats, v2p_map, self.cfg.mode) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, voxel_coords.int(), spatial_shape, self.cfg.batch_size) output = self.input_conv(input_) output = self.unet(output) output = self.output_layer(output) output_feats = output.features[p2v_map.long()] ## F #### semantic segmentation semantic_scores = self.linear(output_feats) # (N, nClass), float semantic_preds = semantic_scores.max(1)[1] # (N), long ---- S data_dict['semantic_scores'] = semantic_scores #### offset pt_offsets_feats = self.offset(output_feats) pt_offsets = self.offset_linear( pt_offsets_feats) # (N, 3), float32 --- O data_dict['pt_offsets'] = pt_offsets if (epoch > self.prepare_epochs): #### get prooposal clusters object_idxs = torch.nonzero(semantic_preds >= 0).view(-1) batch_idxs_ = batch_idxs[object_idxs] batch_offsets_ = utils.get_batch_offsets(batch_idxs_, input_.batch_size) coords_ = coords_float[object_idxs] pt_offsets_ = pt_offsets[object_idxs] semantic_preds_cpu = semantic_preds[object_idxs].int().cpu() #Q idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p( coords_ + pt_offsets_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_shift_meanActive) proposals_idx_shift, proposals_offset_shift = pointgroup_ops.bfs_cluster( semantic_preds_cpu, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre) proposals_idx_shift[:, 1] = object_idxs[ proposals_idx_shift[:, 1].long()].int() # proposals_idx_shift: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # proposals_offset_shift: (nProposal + 1), int #P idx, start_len = pointgroup_ops.ballquery_batch_p( coords_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_meanActive) proposals_idx, proposals_offset = pointgroup_ops.bfs_cluster( semantic_preds_cpu, idx.cpu(), start_len.cpu(), self.cluster_npoint_thre) proposals_idx[:, 1] = object_idxs[proposals_idx[:, 1].long()].int() # proposals_idx: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # proposals_offset: (nProposal + 1), int proposals_idx_shift[:, 0] += (proposals_offset.size(0) - 1) proposals_offset_shift += proposals_offset[-1] proposals_idx = torch.cat((proposals_idx, proposals_idx_shift), dim=0) proposals_offset = torch.cat( (proposals_offset, proposals_offset_shift[1:])) #### proposals voxelization again input_feats, inp_map = self.clusters_voxelization( proposals_idx, proposals_offset, output_feats, coords_float, self.score_fullscale, self.score_scale, self.mode) #### score score = self.score_unet(input_feats) score = self.score_outputlayer(score) score_feats = score.features[inp_map.long()] # (sumNPoint, C) score_feats = pointgroup_ops.roipool( score_feats, proposals_offset.cuda()) # (nProposal, C) scores = self.score_linear(score_feats) # (nProposal, 1) object_feats = self.object_features_linear(score_feats) data_dict['proposal_scores'] = (scores, proposals_idx, proposals_offset) data_dict['proposal_feature'] = object_feats return data_dict
def forward(self, input, input_map, coords, rgb, ori_coords, batch_idxs, batch_offsets, epoch): ''' :param input_map: (N), int, cuda :param coords: (N, 3), float, cuda :param batch_idxs: (N), int, cuda :param batch_offsets: (B + 1), int, cuda ''' ret = {} semantic_scores = [] point_offset_preds = [] voxel_feats = pointgroup_ops.voxelization( input['pt_feats'], input['v2p_map'], input['mode']) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, input['voxel_coords'], input['spatial_shape'], input['batch_size']) output = self.input_conv(input_) output = self.unet(output) output = self.output_layer(output) output_feats = output.features[input_map.long()] output_feats = output_feats.squeeze(dim=0) ### point prediction #### point semantic label prediction semantic_scores.append( self.point_semantic(output_feats)) # (N, nClass), float # point_semantic_preds = semantic_scores point_semantic_preds = semantic_scores[0].max(1)[1] #### point offset prediction point_offset_preds.append( self.point_offset(output_feats)) # (N, 3), float32 point_feats = [] sampled_indexes = [] for sample_indx in range(1, len(batch_offsets)): coords_input = coords[ batch_offsets[sample_indx - 1]:batch_offsets[sample_indx], :].unsqueeze( dim=0) rgb_input = rgb[ batch_offsets[sample_indx - 1]:batch_offsets[sample_indx], :].unsqueeze( dim=0) if not input['test'] and self.center_clustering['use_gt_semantic']: point_semantic_preds = input['semantic_labels'] point_semantic_pred = point_semantic_preds[ batch_offsets[sample_indx - 1]:batch_offsets[sample_indx]] for semantic_idx in point_semantic_pred.unique(): if semantic_idx < 2: continue semantic_point_idx = ( point_semantic_pred == semantic_idx).nonzero().squeeze( dim=1) sampled_index = pointnet2_utils.furthest_point_sample( coords_input[:, semantic_point_idx, :3].contiguous(), self.pointnet_max_npoint).squeeze(dim=0).long() sampled_index = semantic_point_idx[sampled_index] sampled_indexes.append( torch.cat((torch.LongTensor(sampled_index.shape[0], 1).fill_(sample_indx).cuda(), sampled_index.unsqueeze(dim=1)), dim=1)) sampled_coords_input = coords_input[:, sampled_index, :] sampled_rgb_input = rgb_input[:, sampled_index, :] point_feat, _ = self.pointnet_encoder( sampled_coords_input, torch.cat((sampled_rgb_input, sampled_coords_input), dim=2).transpose(1, 2).contiguous()) point_feats.append(point_feat) point_feats = torch.cat(point_feats, dim=0) sampled_indexes = torch.cat(sampled_indexes, dim=0) ### center prediction center_preds = self.center_pred(point_feats) center_semantic_preds = self.center_semantic(point_feats) center_offset_preds = self.center_offset(point_feats) ret['point_semantic_scores'] = semantic_scores ret['point_offset_preds'] = point_offset_preds ret['center_preds'] = (center_preds, sampled_indexes) ret['center_semantic_preds'] = (center_semantic_preds, sampled_indexes) ret['center_offset_preds'] = (center_offset_preds, sampled_indexes) ret['point_features'] = output_feats return ret
def forward(self, input, input_map, coords, rgb, ori_coords, batch_idxs, batch_offsets, epoch): ''' :param input_map: (N), int, cuda :param coords: (N, 3), float, cuda :param batch_idxs: (N), int, cuda :param batch_offsets: (B + 1), int, cuda ''' ret = {} point_semantic_scores = [] point_offset_preds = [] local_point_semantic_scores = [] local_point_offset_preds = [] '''point feature extraction''' # voxelize point cloud, and those voxels are passed to a sparse 3D U-Net # the output of sparse 3D U-Net is then remapped back to points voxel_feats = pointgroup_ops.voxelization( input['pt_feats'], input['v2p_map'], input['mode']) # (M, C), float, cuda input_ = spconv.SparseConvTensor(voxel_feats, input['voxel_coords'], input['spatial_shape'], input['batch_size']) output = self.input_conv(input_) output = self.unet(output) output = self.output_layer(output) output_feats = output.features[input_map.long()] '''point-wise predictions''' point_semantic_scores.append( self.point_semantic(output_feats)) # (N, nClass), float point_semantic_preds = point_semantic_scores[-1].max(1)[1] point_offset_preds.append( self.point_offset(output_feats)) # (N, 3), float32 # cluster proposals with the stable output of the backbone network if (epoch > self.prepare_epochs): if not input['test'] and self.local_proposal['use_gt_semantic']: point_semantic_preds = input['semantic_labels'] '''clustering algorithm''' object_idxs = torch.nonzero(point_semantic_preds > 1).view(-1) batch_idxs_ = batch_idxs[object_idxs] batch_offsets_ = utils.get_batch_offsets(batch_idxs_, input['batch_size']) coords_ = coords[object_idxs] pt_offsets_ = point_offset_preds[-1][object_idxs] semantic_preds_cpu = point_semantic_preds[object_idxs].int().cpu() idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p( coords_ + pt_offsets_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_shift_meanActive) proposals_idx_shift, proposals_offset_shift = pointgroup_ops.bfs_cluster( semantic_preds_cpu, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre) proposals_idx_shift[:, 1] = object_idxs[ proposals_idx_shift[:, 1].long()].int() # proposals_idx_shift: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # proposals_offset_shift: (nProposal + 1), int '''local proposal refinement''' local_proposals_idx = [] local_proposals_offset = [0] # compute proposal centers proposal_center_coords = [] for proposal_index in range(1, len(proposals_offset_shift)): proposal_point_index = proposals_idx_shift[ proposals_offset_shift[proposal_index - 1]: proposals_offset_shift[proposal_index], 1] proposal_center_coords.append( coords[proposal_point_index.long(), :].mean(dim=0, keepdim=True)) proposal_center_coords = torch.cat(proposal_center_coords, dim=0) # select the topk closest proposals for each proposal proposal_dist_mat = euclidean_dist(proposal_center_coords, proposal_center_coords) proposal_dist_mat[range(len(proposal_dist_mat)), range(len(proposal_dist_mat))] = 100 closest_proposals_dist, closest_proposals_index = proposal_dist_mat.topk( k=min(self.local_proposal['topk'], proposal_dist_mat.shape[1]), dim=1, largest=False) valid_closest_proposals_index = closest_proposals_dist < self.local_proposal[ 'dist_th'] # select proposals which are the closest and in the distance threshold for proposal_index in range(1, len(proposals_offset_shift)): local_indexs = [] local_indexs.append(proposals_idx_shift[proposals_offset_shift[ proposal_index - 1]:proposals_offset_shift[proposal_index]]) closest_proposals_ind = closest_proposals_index[ proposal_index - 1, :] for selected_proposal_index in closest_proposals_ind[ valid_closest_proposals_index[proposal_index - 1, :]]: local_index = proposals_idx_shift[ proposals_offset_shift[selected_proposal_index]: proposals_offset_shift[selected_proposal_index + 1]][:, 1].unsqueeze(dim=1) local_indexs.append( torch.cat((torch.LongTensor( local_index.shape[0], 1).fill_(proposal_index - 1).int(), local_index.cpu()), dim=1)) local_proposals_idx.append(torch.cat(local_indexs, dim=0)) local_proposals_offset.append(local_proposals_offset[-1] + local_proposals_idx[-1].shape[0]) local_proposals_idx = torch.cat(local_proposals_idx, dim=0) local_proposals_offset = torch.tensor(local_proposals_offset).int() #### proposals voxelization again input_feats, inp_map = self.clusters_voxelization( local_proposals_idx, local_proposals_offset, output_feats, coords, self.local_proposal['local_proposal_full_scale'], self.local_proposal['local_proposal_scale'], self.mode) #### cluster features if self.local_proposal['reuse_backbone_unet']: proposals = self.unet(input_feats) proposals = self.output_layer(proposals) else: proposals = self.proposal_unet(input_feats) proposals = self.proposal_outputlayer(proposals) proposals_point_features = proposals.features[ inp_map.long()] # (sumNPoint, C) ret['proposals_point_features'] = (proposals_point_features, local_proposals_idx) ### scatter mean point predictions if self.local_proposal['scatter_mean_target'] == 'prediction': refined_point_semantic_score = point_semantic_scores[-1] local_point_semantic_score = self.point_semantic( proposals_point_features) refined_point_semantic_score[:min((local_proposals_idx[:, 1].max() + 1).item(), coords.shape[0]), :] = \ scatter_mean(local_point_semantic_score, local_proposals_idx[:, 1].cuda().long(), dim=0) point_semantic_scores.append(refined_point_semantic_score) point_semantic_preds = refined_point_semantic_score.max(1)[1] refined_point_offset_pred = point_offset_preds[-1] local_point_offset_pred = self.point_offset( proposals_point_features) refined_point_offset_pred[:min((local_proposals_idx[:, 1].max() + 1).item(), coords.shape[0]), :] = \ scatter_mean(local_point_offset_pred, local_proposals_idx[:, 1].cuda().long(), dim=0) point_offset_preds.append(refined_point_offset_pred) ### scatter mean point features elif self.local_proposal['scatter_mean_target'] == 'feature': refined_point_features = torch.zeros_like(output_feats).cuda() refined_point_features[:min((local_proposals_idx[:, 1].max() + 1).item(), coords.shape[0]), :] = \ scatter_mean(proposals_point_features, local_proposals_idx[:, 1].cuda().long(), dim=0) #### filling 0 rows with output_feats filled_index = (refined_point_features == torch.zeros_like( refined_point_features[0, :])).all(dim=1) refined_point_features[filled_index, :] = output_feats[ filled_index, :] ### refined point prediction #### refined point semantic label prediction point_semantic_scores.append( self.point_semantic( refined_point_features)) # (N, nClass), float point_semantic_preds = point_semantic_scores[-1].max(1)[1] #### point offset prediction point_offset_preds.append( self.point_offset( output_feats + refined_point_features)) # (N, 3), float32 elif self.local_proposal['scatter_mean_target'] == False: local_point_semantic_scores.append( self.point_semantic(proposals_point_features)) local_point_offset_preds.append( self.point_offset(proposals_point_features)) ret['local_point_semantic_scores'] = ( local_point_semantic_scores, local_proposals_idx) ret['local_point_offset_preds'] = (local_point_offset_preds, local_proposals_idx) if input['test']: if self.local_proposal['scatter_mean_target'] == False: local_point_semantic_score = self.point_semantic( proposals_point_features) local_point_offset_pred = self.point_offset( proposals_point_features) local_point_semantic_score = scatter_mean( local_point_semantic_score, local_proposals_idx[:, 1].cuda().long(), dim=0) point_semantic_score = point_semantic_scores[-1] point_semantic_score[:local_point_semantic_score.shape[ 0], :] += local_point_semantic_score point_semantic_preds = point_semantic_score.max(1)[1] point_offset_pred = point_offset_preds[-1] local_point_offset_pred = scatter_mean( local_point_offset_pred, local_proposals_idx[:, 1].cuda().long(), dim=0) point_offset_pred[:local_point_offset_pred. shape[0], :] = local_point_offset_pred self.cluster_sets = 'Q' scores, proposals_idx, proposals_offset = self.pointgroup_cluster_algorithm( coords, point_offset_pred, point_semantic_preds, batch_idxs, input['batch_size']) ret['proposal_scores'] = (scores, proposals_idx, proposals_offset) ret['point_semantic_scores'] = point_semantic_scores ret['point_offset_preds'] = point_offset_preds ret['point_features'] = output_feats return ret
def forward(self, input, input_map, coords, rgb, ori_coords, batch_idxs, batch_offsets, epoch): ''' :param input_map: (N), int, cuda :param coords: (N, 3), float, cuda :param batch_idxs: (N), int, cuda :param batch_offsets: (B + 1), int, cuda ''' ret = {} batch_idxs = batch_idxs.squeeze() point_offset_preds = [] point_semantic_scores = [] proposals_confidence_preds = [] proposals_idx_shifts = [] proposals_offset_shifts = [] voxel_feats = pointgroup_ops.voxelization(input['pt_feats'], input['v2p_map'], input['mode']) # (M, C), float, cuda input_ = spconv.SparseConvTensor( voxel_feats, input['voxel_coords'], input['spatial_shape'], input['batch_size'] ) output = self.input_conv(input_) output = self.unet(output) output = self.output_layer(output) output_feats = output.features[input_map.long()] output_feats = output_feats.squeeze(dim=0) ### point prediction #### point semantic label prediction point_semantic_scores.append(self.point_semantic(output_feats)) # (N, nClass), float # point_semantic_preds = semantic_scores point_semantic_preds = point_semantic_scores[-1].max(1)[1] #### point offset prediction point_offset_preds.append(self.point_offset(output_feats)) # (N, 3), float32 point_features = output_feats.clone() if (epoch > self.prepare_epochs): for _ in range(self.proposal_refinement['refine_times']): #### get prooposal clusters object_idxs = torch.nonzero(point_semantic_preds > 1).view(-1) batch_idxs_ = batch_idxs[object_idxs] batch_offsets_ = utils.get_batch_offsets(batch_idxs_, input['batch_size']) coords_ = coords[object_idxs] pt_offsets_ = point_offset_preds[-1][object_idxs] semantic_preds_cpu = point_semantic_preds[object_idxs].int().cpu() idx_shift, start_len_shift = pointgroup_ops.ballquery_batch_p(coords_ + pt_offsets_, batch_idxs_, batch_offsets_, self.cluster_radius, self.cluster_shift_meanActive) proposals_idx_shift, proposals_offset_shift = pointgroup_ops.bfs_cluster(semantic_preds_cpu, idx_shift.cpu(), start_len_shift.cpu(), self.cluster_npoint_thre) proposals_idx_shift[:, 1] = object_idxs[proposals_idx_shift[:, 1].long()].int() # proposals_idx_shift: (sumNPoint, 2), int, dim 0 for cluster_id, dim 1 for corresponding point idxs in N # proposals_offset_shift: (nProposal + 1), int if proposals_idx_shift.shape[0] == 0: continue #### proposals voxelization again input_feats, inp_map = self.clusters_voxelization( proposals_idx_shift, proposals_offset_shift, output_feats, coords, self.proposal_refinement['proposal_refine_full_scale'], self.proposal_refinement['proposal_refine_scale'], self.mode ) #### proposal features proposals = self.proposal_unet(input_feats) proposals = self.proposal_outputlayer(proposals) proposal_feats = proposals.features[inp_map.long()] # (sumNPoint, C) proposal_feats = pointgroup_ops.roipool(proposal_feats, proposals_offset_shift.cuda()) # (nProposal, C) proposals_confidence_preds.append(self.proposal_confidence_linear(proposal_feats)) # (nProposal, 1) proposals_idx_shifts.append(proposals_idx_shift) proposals_offset_shifts.append(proposals_offset_shift) proposal_pts_idxs = proposals_idx_shift[proposals_offset_shift[:-1].long()][:, 1].cuda() if len(proposal_pts_idxs) == 0: continue proposal_batch_idxs = proposal_pts_idxs.clone() for _batch_idx in range(len(batch_offsets_) - 1, 0, -1): proposal_batch_idxs[proposal_pts_idxs < batch_offsets_[_batch_idx]] = _batch_idx refined_point_features = [] for _batch_idx in range(1, len(batch_offsets)): key_input = proposal_feats[proposal_batch_idxs == _batch_idx, :].unsqueeze(dim=0) query_input = output_feats[batch_offsets[_batch_idx - 1]:batch_offsets[_batch_idx], :].unsqueeze(dim=0).permute(0, 2, 1) # query_input = output_feats[batch_offsets[_batch_idx - 1]:batch_offsets[_batch_idx-1]+100, # :].unsqueeze(dim=0).permute(0, 2, 1) point_refined_feature, _ = self.proposal_transformer( src=key_input, query_embed=query_input, ) refined_point_features.append(point_refined_feature.squeeze(dim=0)) refined_point_features = torch.cat(refined_point_features, dim=0) assert refined_point_features.shape[0] == point_features.shape[ 0], 'point wise features have wrong point numbers' point_features = refined_point_features.clone() ### refined point prediction #### refined point semantic label prediction point_semantic_scores.append(self.point_semantic(refined_point_features)) # (N, nClass), float point_semantic_preds = point_semantic_scores[-1].max(1)[1] #### point offset prediction point_offset_preds.append(self.point_offset(refined_point_features)) # (N, 3), float32 if (epoch == self.test_epoch) and input['test']: self.cluster_sets = 'Q' scores, proposals_idx, proposals_offset = self.pointgroup_cluster_algorithm( coords, point_offset_preds[-1], point_semantic_preds, batch_idxs, input['batch_size'] ) ret['proposal_scores'] = (scores, proposals_idx, proposals_offset) ret['proposal_confidences'] = (proposals_confidence_preds, proposals_idx_shifts, proposals_offset_shifts) ret['point_semantic_scores'] = point_semantic_scores ret['point_offset_preds'] = point_offset_preds return ret