# 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(