# apply the observation mask view_s,mask_s,_ = util.apply_mask(complete_s.clone(),args.maskMethod) view_t,mask_t,_ = util.apply_mask(complete_t.clone(),args.maskMethod) mask_s=torch_op.npy(mask_s[0,:,:,:]).transpose(1,2,0) mask_t=torch_op.npy(mask_t[0,:,:,:]).transpose(1,2,0) # append mask for valid data tpmask = (view_s[:,6:7,:,:]!=0).float().cuda() view_s=torch.cat((view_s,tpmask),1) tpmask = (view_t[:,6:7,:,:]!=0).float().cuda() view_t=torch.cat((view_t,tpmask),1) # warp the second scan using current transformation estimation view_t2s=torch_op.v(util.warping(torch_op.npy(view_t),np.linalg.inv(R_hat),args.dataset)) view_s2t=torch_op.v(util.warping(torch_op.npy(view_s),R_hat,args.dataset)) # append the warped scans #view0 = torch.cat((view_s,view_t2s),1) #view1 = torch.cat((view_t,view_s2t),1) view0 = torch.cat((view_s,view_s),1) view1 = torch.cat((view_t,view_t),1) # generate complete scans f=net(torch.cat((view0,view1))) f0=f[0:1,:,:,:] f1=f[1:2,:,:,:] data_sc,data_tc={},{} complete_normal_s.append(torch_op.npy(f0[0,3:6,:,:]))
def RelativePoseEstimationViaCompletion(net, data_s, data_t, args): """ The main algorithm: Given two set of scans, alternate between scan completion and pairwise matching args need to contain: snumclass: number of semantic class featureDim: feature dimension outputType: ['rgb':color,'d':depth,'n':normal,'s':semantic,'f':feature] maskMethod: ['second'] alterStep: dataset: para: """ EPS = 1e-12 args.idx_f_start = 0 if 'rgb' in args.outputType: args.idx_f_start += 3 if 'n' in args.outputType: args.idx_f_start += 3 if 'd' in args.outputType: args.idx_f_start += 1 if 's' in args.outputType: args.idx_f_start += args.snumclass if 'f' in args.outputType: args.idx_f_end = args.idx_f_start + args.featureDim with torch.set_grad_enabled(False): R_hat=np.eye(4) # get the complete scans complete_s=torch.cat((torch_op.v(data_s['rgb']),torch_op.v(data_s['norm']),torch_op.v(data_s['depth']).unsqueeze(2)),2).permute(2,0,1).unsqueeze(0) complete_t=torch.cat((torch_op.v(data_t['rgb']),torch_op.v(data_t['norm']),torch_op.v(data_t['depth']).unsqueeze(2)),2).permute(2,0,1).unsqueeze(0) # apply the observation mask view_s,mask_s,_ = util.apply_mask(complete_s.clone(),args.maskMethod) view_t,mask_t,_ = util.apply_mask(complete_t.clone(),args.maskMethod) mask_s=torch_op.npy(mask_s[0,:,:,:]).transpose(1,2,0) mask_t=torch_op.npy(mask_t[0,:,:,:]).transpose(1,2,0) # append mask for valid data tpmask = (view_s[:,6:7,:,:]!=0).float().cuda() view_s=torch.cat((view_s,tpmask),1) tpmask = (view_t[:,6:7,:,:]!=0).float().cuda() view_t=torch.cat((view_t,tpmask),1) for alter_ in range(args.alterStep): # warp the second scan using current transformation estimation view_t2s=torch_op.v(util.warping(torch_op.npy(view_t),np.linalg.inv(R_hat),args.dataset)) view_s2t=torch_op.v(util.warping(torch_op.npy(view_s),R_hat,args.dataset)) # append the warped scans view0 = torch.cat((view_s,view_t2s),1) view1 = torch.cat((view_t,view_s2t),1) # generate complete scans f=net(torch.cat((view0,view1))) f0=f[0:1,:,:,:] f1=f[1:2,:,:,:] data_sc,data_tc={},{} # replace the observed region with gt depth/normal data_sc['normal'] = (1-mask_s)*torch_op.npy(f0[0,3:6,:,:]).transpose(1,2,0)+mask_s*data_s['norm'] data_tc['normal'] = (1-mask_t)*torch_op.npy(f1[0,3:6,:,:]).transpose(1,2,0)+mask_t*data_t['norm'] data_sc['normal']/= (np.linalg.norm(data_sc['normal'],axis=2,keepdims=True)+EPS) data_tc['normal']/= (np.linalg.norm(data_tc['normal'],axis=2,keepdims=True)+EPS) data_sc['depth'] = (1-mask_s[:,:,0])*torch_op.npy(f0[0,6,:,:])+mask_s[:,:,0]*data_s['depth'] data_tc['depth'] = (1-mask_t[:,:,0])*torch_op.npy(f1[0,6,:,:])+mask_t[:,:,0]*data_t['depth'] data_sc['obs_mask'] = mask_s.copy() data_tc['obs_mask'] = mask_t.copy() data_sc['rgb'] = (mask_s*data_s['rgb']*255).astype('uint8') data_tc['rgb'] = (mask_t*data_t['rgb']*255).astype('uint8') # for scannet, we use the original size rgb image(480x640) to extract sift keypoint if 'scannet' in args.dataset: data_sc['rgb_full'] = (data_s['rgb_full']*255).astype('uint8') data_tc['rgb_full'] = (data_t['rgb_full']*255).astype('uint8') data_sc['depth_full'] = data_s['depth_full'] data_tc['depth_full'] = data_t['depth_full'] # extract feature maps f0_feat=f0[:,args.idx_f_start:args.idx_f_end,:,:] f1_feat=f1[:,args.idx_f_start:args.idx_f_end,:,:] data_sc['feat']=f0_feat.squeeze(0) data_tc['feat']=f1_feat.squeeze(0) para_this = copy.copy(args.para) para_this.sigmaAngle1 = para_this.sigmaAngle1[alter_] para_this.sigmaAngle2 = para_this.sigmaAngle2[alter_] para_this.sigmaDist = para_this.sigmaDist[alter_] para_this.sigmaFeat = para_this.sigmaFeat[alter_] # run relative pose module to get next estimate R_hat = RelativePoseEstimation(data_sc,data_tc,para_this,args.dataset,args.representation,doCompletion=args.completion,maskMethod=args.maskMethod,index=None) return R_hat
mask_s = torch_op.npy(mask_s[0, :, :, :]).transpose( 1, 2, 0) mask_t = torch_op.npy(mask_t[0, :, :, :]).transpose( 1, 2, 0) # append mask for valid data tpmask = (view_s[:, 6:7, :, :] != 0).float().cuda() view_s = torch.cat((view_s, tpmask), 1) tpmask = (view_t[:, 6:7, :, :] != 0).float().cuda() view_t = torch.cat((view_t, tpmask), 1) for alter_ in range(args.alterStep): # warp the second scan using current transformation estimation view_t2s = torch_op.v( util.warping(torch_op.npy(view_t), np.linalg.inv(R_hat), args.dataList)) view_s2t = torch_op.v( util.warping(torch_op.npy(view_s), R_hat, args.dataList)) # append the warped scans view0 = torch.cat((view_s, view_t2s), 1) view1 = torch.cat((view_t, view_s2t), 1) # generate complete scans f = net(torch.cat((view0, view1))) f0 = f[0:1, :, :, :] f1 = f[1:2, :, :, :] data_sc, data_tc = {}, {} # replace the observed region with observed depth/normal data_sc['normal'] = (1 - mask_s) * torch_op.npy(