def test(modelin=args.model,outfile=args.out,optimize=args.opt,ft=args.ft): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1,feature_transform=ft) sfm_net = PointNet(n=199,feature_transform=ft) if modelin != "": calib_path = os.path.join('model','calib_' + modelin) sfm_path = os.path.join('model','sfm_' + modelin) calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.eval() sfm_net.eval() # mean shape and eigenvectors for 3dmm M = 100 data3dmm = dataloader.SyntheticLoader() mu_lm = torch.from_numpy(data3dmm.mu_lm).float().detach() mu_lm[:,2] = mu_lm[:,2]*-1 lm_eigenvec = torch.from_numpy(data3dmm.lm_eigenvec).float().detach() sigma = torch.from_numpy(data3dmm.sigma).float().detach() sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) # sample from f testing set allerror_2d = [] allerror_3d = [] allerror_rel3d = [] allerror_relf = [] all_f = [] all_fpred = [] all_depth = [] out_shape = [] out_f = [] seterror_3d = [] seterror_rel3d = [] seterror_relf = [] seterror_2d = [] f_vals = [i*100 for i in range(4,15)] # set random seed for reproducibility of test set np.random.seed(0) torch.manual_seed(0) for f_test in f_vals: # create dataloader loader = dataloader.TestLoader(f_test) f_pred = [] shape_pred = [] error_2d = [] error_3d = [] error_rel3d = [] error_relf = [] M = 100; N = 68; batch_size = 1; for j,data in enumerate(loader): if j >= 10: break # load the data x_cam_gt = data['x_cam_gt'] shape_gt = data['x_w_gt'] fgt = data['f_gt'] x_img = data['x_img'] x_img_gt = data['x_img_gt'] depth = torch.norm(x_cam_gt.mean(2),dim=1) all_depth.append(depth.numpy()) all_f.append(fgt.numpy()[0]) ptsI = x_img.reshape((M,N,2)).permute(0,2,1) x = x_img.unsqueeze(0).permute(0,2,1) # run the model f = calib_net(x) + 300 betas = sfm_net(x) betas = betas.squeeze(0).unsqueeze(-1) shape = mu_lm + torch.mm(lm_eigenvec,betas).squeeze().view(N,3) shape = shape - shape.mean(0).unsqueeze(0) # get motion measurement guess K = torch.zeros((3,3)).float() K[0,0] = f K[1,1] = f K[2,2] = 1 km,c_w,scaled_betas,alphas = util.EPnP(ptsI,shape,K) _, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI) error_time = util.getTimeConsistency(shape,R,T) if error_time > 10: mode='walk' else: mode='still' # apply dual optimization if optimize: calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) shape,K,R,T = dualoptimization(x,calib_net,sfm_net,shape_gt=shape_gt,fgt=fgt,mode=mode) f = K[0,0].detach() else: K = torch.zeros(3,3).float() K[0,0] = f K[1,1] = f K[2,2] = 1 km,c_w,scaled_betas,alphas = util.EPnP(ptsI,shape,K) Xc, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI) # get errors reproj_errors2 = util.getReprojError2(ptsI,shape,R,T,K,show=False) reproj_errors3 = torch.norm(shape_gt - shape,dim=1).mean() rel_errors = util.getRelReprojError3(x_cam_gt,shape,R,T) reproj_error = reproj_errors2.mean() reconstruction_error = reproj_errors3.mean() rel_error = rel_errors.mean() f_error = torch.abs(fgt - f) / fgt # save final prediction f_pred.append(f.detach().cpu().item()) shape_pred.append(shape.detach().cpu().numpy()) all_fpred.append(f.detach().data.numpy()) allerror_3d.append(reproj_error.data.numpy()) allerror_2d.append(reconstruction_error.data.numpy()) allerror_rel3d.append(rel_error.data.numpy()) error_2d.append(reproj_error.cpu().data.item()) error_3d.append(reconstruction_error.cpu().data.item()) error_rel3d.append(rel_error.cpu().data.item()) error_relf.append(f_error.cpu().data.item()) print(f"f/sequence: {f_test}/{j} | f/fgt: {f.item():.3f}/{fgt.item():.3f} | f_error_rel: {f_error.item():.4f} | rmse: {reconstruction_error.item():.4f} | rel rmse: {rel_error.item():.4f} | 2d error: {reproj_error.item():.4f}") avg_2d = np.mean(error_2d) avg_rel3d = np.mean(error_rel3d) avg_3d = np.mean(error_3d) avg_relf = np.mean(error_relf) seterror_2d.append(avg_2d) seterror_3d.append(avg_3d) seterror_rel3d.append(avg_rel3d) seterror_relf.append(avg_relf) out_f.append(np.stack(f_pred)) out_shape.append(np.concatenate(shape_pred,axis=0)) print(f"f_error_rel: {avg_relf:.4f} | rel rmse: {avg_rel3d:.4f} | 2d error: {avg_2d:.4f} | rmse: {avg_3d:.4f} |") # save output out_shape = np.stack(out_shape) out_f = np.stack(out_f) all_f = np.stack(all_f).flatten() all_fpred = np.stack(all_fpred).flatten() all_d = np.stack(all_depth).flatten() allerror_2d = np.stack(allerror_2d).flatten() allerror_rel3d = np.stack(allerror_rel3d).flatten() matdata = {} matdata['fvals'] = np.array(f_vals) matdata['all_f'] = np.array(all_f) matdata['all_fpred'] = np.array(all_fpred) matdata['all_d'] = np.array(all_depth) matdata['error_2d'] = allerror_2d matdata['error_3d'] = allerror_3d matdata['error_rel3d'] = allerror_rel3d matdata['seterror_2d'] = np.array(seterror_2d) matdata['seterror_3d'] = np.array(seterror_3d) matdata['seterror_rel3d'] = np.array(seterror_rel3d) matdata['seterror_relf'] = np.array(seterror_relf) matdata['shape'] = np.stack(out_shape) matdata['f'] = np.stack(out_f) scipy.io.savemat(outfile,matdata) print(f"MEAN seterror_2d: {np.mean(seterror_2d)}") print(f"MEAN seterror_3d: {np.mean(seterror_3d)}") print(f"MEAN seterror_rel3d: {np.mean(seterror_rel3d)}") print(f"MEAN seterror_relf: {np.mean(seterror_relf)}") return np.mean(seterror_relf)
def train(modelin=args.model, modelout=args.out,device=args.device,opt=args.opt,ft=args.ft): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1,feature_transform=ft) sfm_net = PointNet(n=199,feature_transform=ft) if modelin != "": calib_path = os.path.join('model','calib_' + modelin) sfm_path = os.path.join('model','sfm_' + modelin) pretrained1 = torch.load(calib_path) pretrained2 = torch.load(sfm_path) calib_dict = calib_net.state_dict() sfm_dict = sfm_net.state_dict() pretrained1 = {k: v for k,v in pretrained1.items() if k in calib_dict} pretrained2 = {k: v for k,v in pretrained2.items() if k in sfm_dict} calib_dict.update(pretrained1) sfm_dict.update(pretrained2) calib_net.load_state_dict(pretrained1) sfm_net.load_state_dict(pretrained2) calib_net.to(device=device) sfm_net.to(device=device) opt1 = torch.optim.Adam(calib_net.parameters(),lr=1e-3) opt2 = torch.optim.Adam(sfm_net.parameters(),lr=1e-3) # dataloader loader = dataloader.SyntheticLoader() batch_size = 100 M = loader.M N = loader.N # mean shape and eigenvectors for 3dmm mu_lm = torch.from_numpy(loader.mu_lm).float()#.to(device=device) mu_lm[:,2] = mu_lm[:,2] * -1 mu_lm = torch.stack(300 * [mu_lm.to(device=device)]) shape = mu_lm lm_eigenvec = torch.from_numpy(loader.lm_eigenvec).float().to(device=device) sigma = torch.from_numpy(loader.sigma).float().detach().to(device=device) sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) lm_eigenvec = torch.stack(300 * [lm_eigenvec]) # main training loop best = 10000 for epoch in itertools.count(): for i in range(len(loader)): if i < 3: continue v1 = loader[i] v2 = loader[i-1] v3 = loader[i-2] batch = stackVideos([v1,v2,v3],100,68,device=device) # get the input and gt values alpha_gt = batch['alpha'] x_cam_gt = batch['x_cam_gt'] shape_gt = batch['x_w_gt'] fgt = batch['f_gt'] x = batch['x_img'] M = x.shape[0] N = x.shape[-1] # calibration f = torch.squeeze(calib_net(x) + 300) K = torch.zeros((M,3,3)).float().to(device=device) K[:,0,0] = f K[:,1,1] = f K[:,2,2] = 1 # sfm alpha = sfm_net(x) alpha = alpha.unsqueeze(-1) shape = mu_lm + torch.bmm(lm_eigenvec,alpha).squeeze().view(M,N,3) shape[0:100] = shape[0:100] - shape[0:100].mean(1).unsqueeze(1) shape[100:200] = shape[100:200] - shape[100:200].mean(1).unsqueeze(1) shape[200:300] = shape[200:300] - shape[200:300].mean(1).unsqueeze(1) opt1.zero_grad() opt2.zero_grad() f1_error = torch.mean(torch.abs(f[0:100] - fgt[0])) f2_error = torch.mean(torch.abs(f[100:200] - fgt[1])) f3_error = torch.mean(torch.abs(f[200:300] - fgt[2])) #a1_error = torch.mean(torch.abs(alpha[0:100] - alpha_gt[0])) #a2_error = torch.mean(torch.abs(alpha[100:200] - alpha_gt[1])) #a3_error = torch.mean(torch.abs(alpha[200:300] - alpha_gt[2])) s1_error = torch.mean(torch.abs(shape[0:100] - shape_gt[0].unsqueeze(0))) s2_error = torch.mean(torch.abs(shape[100:200] - shape_gt[1].unsqueeze(0))) s3_error = torch.mean(torch.abs(shape[200:300] - shape_gt[2].unsqueeze(0))) ferror = f1_error + f2_error + f3_error #aerror = a1_error + a2_error + a3_error serror = s1_error + s2_error + s3_error #f_error = torch.mean(torch.abs(f - fgt)) #error3d = torch.mean(torch.norm(shape - shape_gt,dim=2)) #error = ferror + aerror error = ferror + serror error.backward() opt1.step() opt2.step() print(f"iter: {i} | best: {best:.2f} | f_error: {ferror.item():.3f} | serror: {serror.item():.3f} ") if i == 1000: break # save model and increment weight decay torch.save(sfm_net.state_dict(), os.path.join('model','sfm_model.pt')) torch.save(calib_net.state_dict(), os.path.join('model','calib_model.pt')) ferror = test(modelin='model.pt',outfile=args.out,optimize=False) if ferror < best: best = ferror print("saving!") torch.save(sfm_net.state_dict(), os.path.join('model','sfm_'+modelout)) torch.save(calib_net.state_dict(), os.path.join('model','calib_'+modelout)) sfm_net.train() calib_net.train()
def testReal(modelin=args.model,outfile=args.out,optimize=args.opt,db=args.db): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1) sfm_net = PointNet(n=199) if modelin != "": calib_path = os.path.join('model','calib_' + modelin) sfm_path = os.path.join('model','sfm_' + modelin) calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.eval() sfm_net.eval() # mean shape and eigenvectors for 3dmm data3dmm = dataloader.SyntheticLoader() mu_lm = torch.from_numpy(data3dmm.mu_lm).float().detach() mu_lm[:,2] = mu_lm[:,2]*-1 lm_eigenvec = torch.from_numpy(data3dmm.lm_eigenvec).float().detach() sigma = torch.from_numpy(data3dmm.sigma).float().detach() sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) # define loader loader = getLoader(db) f_pred = [] shape_pred = [] error_2d = [] error_relf = [] error_rel3d = [] for sub in range(len(loader)): batch = loader[sub] x_cam_gt = batch['x_cam_gt'] fgt = batch['f_gt'] x_img = batch['x_img'] x_img_gt = batch['x_img_gt'] M = x_img_gt.shape[0] N = x_img_gt.shape[-1] ptsI = x_img.reshape((M,N,2)).permute(0,2,1) x = x_img.unsqueeze(0).permute(0,2,1) # run the model f = calib_net(x) + 300 betas = sfm_net(x) betas = betas.squeeze(0).unsqueeze(-1) shape = mu_lm + torch.mm(lm_eigenvec,betas).squeeze().view(N,3) shape = shape - shape.mean(0).unsqueeze(0) # get motion measurement guess K = torch.zeros((3,3)).float() K[0,0] = f K[1,1] = f K[2,2] = 1 km,c_w,scaled_betas,alphas = util.EPnP(ptsI,shape,K) _, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI) error_time = util.getTimeConsistency(shape,R,T) if error_time > 20: mode='walk' else: mode='still' # adjust number of landmarks M = x_img_gt.shape[0] N = x_img_gt.shape[-1] # additional optimization on initial solution if optimize: calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) if db == 'biwi': shape_gt = batch['x_w_gt'] shape,K,R,T = dualoptimization(x,calib_net,sfm_net,shape_gt=shape_gt,fgt=fgt,M=M,N=N,mode=mode,db='biwi') else: shape,K,R,T = dualoptimization(x,calib_net,sfm_net,fgt=fgt,M=M,N=N,mode=mode) f = K[0,0].detach() else: K = torch.zeros(3,3).float() K[0,0] = f K[1,1] = f K[2,2] = 1 km,c_w,scaled_betas,alphas = util.EPnP(ptsI,shape,K) Xc, R, T, mask = util.optimizeGN(km,c_w,scaled_betas,alphas,shape,ptsI) # get errors reproj_errors2 = util.getReprojError2(ptsI,shape,R,T,K) rel_errors = util.getRelReprojError3(x_cam_gt,shape,R,T) reproj_error = reproj_errors2.mean() rel_error = rel_errors.mean() f_error = torch.abs(fgt - f) / fgt # save final prediction f_pred.append(f.detach().cpu().item()) shape_pred.append(shape.detach().cpu().numpy()) error_2d.append(reproj_error.cpu().data.item()) error_rel3d.append(rel_error.cpu().data.item()) error_relf.append(f_error.cpu().data.item()) print(f" f/fgt: {f.item():.3f}/{fgt.item():.3f} | f_error_rel: {f_error.item():.4f} | rel rmse: {rel_error.item():.4f} | 2d error: {reproj_error.item():.4f}") #end for # prepare output file out_shape = np.stack(shape_pred) out_f = np.stack(f_pred) matdata = {} matdata['shape'] = np.stack(out_shape) matdata['f'] = np.stack(out_f) matdata['error_2d'] = np.array(error_2d) matdata['error_rel3d'] = np.array(error_rel3d) matdata['error_relf'] = np.array(error_relf) scipy.io.savemat(outfile,matdata) print(f"MEAN seterror_2d: {np.mean(error_2d)}") print(f"MEAN seterror_rel3d: {np.mean(error_rel3d)}") print(f"MEAN seterror_relf: {np.mean(error_relf)}")
def train(modelin=args.model, modelout=args.out, device=args.device, opt=args.opt): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1) sfm_net = PointNet(n=199) if modelin != "": calib_path = os.path.join('model', 'calib_' + modelin) sfm_path = os.path.join('model', 'sfm_' + modelin) pretrained1 = torch.load(calib_path) pretrained2 = torch.load(sfm_path) calib_dict = calib_net.state_dict() sfm_dict = sfm_net.state_dict() pretrained1 = {k: v for k, v in pretrained1.items() if k in calib_dict} pretrained2 = {k: v for k, v in pretrained2.items() if k in sfm_dict} calib_dict.update(pretrained1) sfm_dict.update(pretrained2) calib_net.load_state_dict(pretrained1) sfm_net.load_state_dict(pretrained2) calib_net.to(device=device) sfm_net.to(device=device) opt1 = torch.optim.Adam(calib_net.parameters(), lr=1e-4) opt2 = torch.optim.Adam(sfm_net.parameters(), lr=1e-2) # dataloader data = dataloader.Data() loader = data.batchloader batch_size = data.batchsize # mean shape and eigenvectors for 3dmm mu_lm = torch.from_numpy(data.mu_lm).float() #.to(device=device) mu_lm[:, 2] = mu_lm[:, 2] * -1 mu_lm = torch.stack(batch_size * [mu_lm.to(device=device)]) shape = mu_lm lm_eigenvec = torch.from_numpy(data.lm_eigenvec).float().to(device=device) lm_eigenvec = torch.stack(batch_size * [lm_eigenvec]) M = data.M N = data.N # main training loop best = 10000 for epoch in itertools.count(): for j, batch in enumerate(loader): # get the input and gt values x_cam_gt = batch['x_cam_gt'].to(device=device) shape_gt = batch['x_w_gt'].to(device=device) fgt = batch['f_gt'].to(device=device) x_img = batch['x_img'].to(device=device) #beta_gt = batch['beta_gt'].to(device=device) #x_img_norm = batch['x_img_norm'] x_img_gt = batch['x_img_gt'].to(device=device).permute(0, 2, 1, 3) batch_size = fgt.shape[0] one = torch.ones(batch_size, M * N, 1).to(device=device) x_img_one = torch.cat([x_img, one], dim=2) x_cam_pt = x_cam_gt.permute(0, 1, 3, 2).reshape(batch_size, 6800, 3) x = x_img.permute(0, 2, 1) #x = x_img.permute(0,2,1).reshape(batch_size,2,M,N) ptsI = x_img_one.reshape(batch_size, M, N, 3).permute(0, 1, 3, 2)[:, :, :2, :] # if just optimizing if not opt: # calibration f = calib_net(x) + 300 K = torch.zeros((batch_size, 3, 3)).float().to(device=device) K[:, 0, 0] = f.squeeze() K[:, 1, 1] = f.squeeze() K[:, 2, 2] = 1 # sfm betas = sfm_net(x) betas = betas.unsqueeze(-1) shape = mu_lm + torch.bmm(lm_eigenvec, betas).squeeze().view( batch_size, N, 3) opt1.zero_grad() opt2.zero_grad() f_error = torch.mean(torch.abs(f - fgt)) #error2d = torch.mean(torch.abs(pred - x_img_gt)) error3d = torch.mean(torch.abs(shape - shape_gt)) error = f_error + error3d error.backward() opt1.step() opt2.step() print( f"{best:.2f} | f_error: {f_error.item():.3f} | error3d: {error3d.item():.3f} | f/fgt: {f[0].item():.1f}/{fgt[0].item():.1f} | f/fgt: {f[1].item():.1f}/{fgt[1].item():.1f} | f/fgt: {f[2].item():.1f}/{fgt[2].item():.1f} | f/fgt: {f[3].item():.1f}/{fgt[3].item():.1f} " ) continue # save model and increment weight decay torch.save(sfm_net.state_dict(), os.path.join('model', 'sfm_model.pt')) torch.save(calib_net.state_dict(), os.path.join('model', 'calib_model.pt')) ferror = test(modelin='model.pt', outfile=args.out, optimize=False) if ferror < best: best = ferror print("saving!") torch.save(sfm_net.state_dict(), os.path.join('model', 'sfm_' + modelout)) torch.save(calib_net.state_dict(), os.path.join('model', 'calib_' + modelout)) sfm_net.train() calib_net.train()
def test_sfm(modelin=args.model, outfile=args.out, optimize=args.opt): # define model, dataloader, 3dmm sfm_net = PointNet(n=199) calib_net = PointNet(n=1) if modelin != "": calib_path = os.path.join('model', 'calib_' + modelin) sfm_path = os.path.join('model', 'sfm_' + modelin) calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.eval() sfm_net.eval() # mean shape and eigenvectors for 3dmm M = 100 data3dmm = dataloader.SyntheticLoader() mu_lm = torch.from_numpy(data3dmm.mu_lm).float().detach() mu_lm[:, 2] = mu_lm[:, 2] * -1 lm_eigenvec = torch.from_numpy(data3dmm.lm_eigenvec).float().detach() sigma = torch.from_numpy(data3dmm.sigma).float().detach() sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) # sample from f testing set allerror_2d = [] allerror_3d = [] allerror_rel3d = [] allerror_relf = [] all_f = [] all_fpred = [] all_depth = [] out_shape = [] out_f = [] seterror_3d = [] seterror_rel3d = [] seterror_relf = [] seterror_2d = [] f_vals = [i * 100 for i in range(4, 15)] for f_test in f_vals: # create dataloader #f_test = 1000 loader = dataloader.TestLoader(f_test, addnoise=True) f_pred = [] shape_pred = [] error_2d = [] error_3d = [] error_rel3d = [] error_relf = [] loader.M = 100 loader.N = 68 batch_size = 1 M = loader.M N = loader.N #training_pred = np.zeros((5,M,68,3)) #training_gt = np.zeros((5,M,68,3)) for j, data in enumerate(loader): if j == 1: break # load the data x_cam_gt = data['x_cam_gt'] shape_gt = data['x_w_gt'] fgt = data['f_gt'] x_img = data['x_img'] x_img_gt = data['x_img_gt'] T_gt = data['T_gt'] all_depth.append(np.mean(T_gt[:, 2])) all_f.append(fgt.numpy()[0]) ptsI = x_img.reshape((M, N, 2)).permute(0, 2, 1) x = x_img.unsqueeze(0).permute(0, 2, 1) # test sfm calibration #calib_net.load_state_dict(torch.load(calib_path)) opt2 = torch.optim.Adam(sfm_net.parameters(), lr=1) sfm_net.eval() trainfc(sfm_net) f = fgt #v = pptk.viewer(mu_lm.cpu().numpy()) l_init = 100000 for iter in itertools.count(): opt2.zero_grad() # shape prediction betas = sfm_net(ptsI).mean(0).unsqueeze(1) shape = mu_lm + torch.mm(lm_eigenvec, betas).squeeze().view( N, 3) shape = shape - shape.mean(0).unsqueeze(0) rmse = torch.norm(shape_gt - shape, dim=1).mean().detach() K = torch.zeros((3, 3)).float() K[0, 0] = f K[1, 1] = f K[2, 2] = 1 km, c_w, scaled_betas, alphas = util.EPnP(ptsI, shape, K) Xc, R, T, mask = util.optimizeGN(km, c_w, scaled_betas, alphas, shape, ptsI) error2d = util.getReprojError2(ptsI, shape, R, T, K, show=False, loss='l2') error_time = util.getTimeConsistency(shape, R, T) loss = error2d.mean() + 0.01 * error_time loss.backward() opt2.step() print( f"iter: {iter} | error: {loss.item():.3f} | error2d: {error2d.mean().item():.3f} | rmse: {rmse.item():.3f} " ) if iter >= 10 and l_init - loss < 0.0001: break print(l_init - loss) l_init = loss #training_pred[j,iter,:,:] = shape.detach().cpu() #training_gt[j,iter,:,:] = shape_gt.detach().cpu() all_fpred.append(f.detach().numpy()[0]) # get errors reproj_errors2 = util.getReprojError2(ptsI, shape, R, T, K, show=False) reproj_errors3 = torch.norm(shape_gt - shape, dim=1).mean() rel_errors = util.getRelReprojError3(x_cam_gt, shape, R, T) reproj_error = reproj_errors2.mean() reconstruction_error = reproj_errors3.mean() rel_error = rel_errors.mean() f_error = torch.abs(fgt - f) / fgt # save final prediction f_pred.append(f.detach().cpu().item()) shape_pred.append(shape.detach().cpu().numpy()) allerror_3d.append(reproj_error.data.numpy()) allerror_2d.append(reconstruction_error.data.numpy()) allerror_rel3d.append(rel_error.data.numpy()) error_2d.append(reproj_error.cpu().data.item()) error_3d.append(reconstruction_error.cpu().data.item()) error_rel3d.append(rel_error.cpu().data.item()) error_relf.append(f_error.cpu().data.item()) print( f"f/sequence: {f_test}/{j} | f/fgt: {f[0].item():.3f}/{fgt.item():.3f} | f_error_rel: {f_error.item():.4f} | rmse: {reconstruction_error.item():.4f} | rel rmse: {rel_error.item():.4f} | 2d error: {reproj_error.item():.4f}" ) avg_2d = np.mean(error_2d) avg_rel3d = np.mean(error_rel3d) avg_3d = np.mean(error_3d) avg_relf = np.mean(error_relf) seterror_2d.append(avg_2d) seterror_3d.append(avg_3d) seterror_rel3d.append(avg_rel3d) seterror_relf.append(avg_relf) out_f.append(np.stack(f_pred)) out_shape.append(np.stack(shape_pred, axis=0)) print( f"f_error_rel: {avg_relf:.4f} | rel rmse: {avg_rel3d:.4f} | 2d error: {reproj_error.item():.4f} | rmse: {avg_3d:.4f} |" ) out_f = np.stack(out_f) all_f = np.stack(all_f).flatten() all_fpred = np.stack(all_fpred).flatten() all_d = np.stack(all_depth).flatten() allerror_2d = np.stack(allerror_2d).flatten() allerror_3d = np.stack(allerror_3d).flatten() allerror_rel3d = np.stack(allerror_rel3d).flatten() matdata = {} #matdata['training_pred'] = training_pred #matdata['training_gt'] = training_gt matdata['fvals'] = np.array(f_vals) matdata['all_f'] = np.array(all_f) matdata['all_fpred'] = np.array(all_fpred) matdata['all_d'] = np.array(all_depth) matdata['error_2d'] = allerror_2d matdata['error_3d'] = allerror_3d matdata['error_rel3d'] = allerror_rel3d matdata['seterror_2d'] = np.array(seterror_2d) matdata['seterror_3d'] = np.array(seterror_3d) matdata['seterror_rel3d'] = np.array(seterror_rel3d) matdata['seterror_relf'] = np.array(seterror_relf) matdata['f'] = np.stack(out_f) scipy.io.savemat(outfile, matdata) print(f"MEAN seterror_2d: {np.mean(seterror_2d)}") print(f"MEAN seterror_3d: {np.mean(seterror_3d)}") print(f"MEAN seterror_rel3d: {np.mean(seterror_rel3d)}") print(f"MEAN seterror_relf: {np.mean(seterror_relf)}")
def train(modelin=args.model, modelout=args.out, device=args.device, opt=args.opt, ft=args.ft): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1, feature_transform=ft) sfm_net = PointNet(n=199, feature_transform=ft) if modelin != "": calib_path = os.path.join('model', 'calib_' + modelin) sfm_path = os.path.join('model', 'sfm_' + modelin) pretrained1 = torch.load(calib_path) pretrained2 = torch.load(sfm_path) calib_dict = calib_net.state_dict() sfm_dict = sfm_net.state_dict() pretrained1 = {k: v for k, v in pretrained1.items() if k in calib_dict} pretrained2 = {k: v for k, v in pretrained2.items() if k in sfm_dict} calib_dict.update(pretrained1) sfm_dict.update(pretrained2) calib_net.load_state_dict(pretrained1) sfm_net.load_state_dict(pretrained2) calib_net.to(device=device) sfm_net.to(device=device) opt1 = torch.optim.Adam(calib_net.parameters(), lr=1e-3) opt2 = torch.optim.Adam(sfm_net.parameters(), lr=1e-3) # dataloader loader = dataloader.SyntheticLoader() batch_size = 100 M = loader.M N = loader.N # mean shape and eigenvectors for 3dmm mu_lm = torch.from_numpy(loader.mu_lm).float() #.to(device=device) mu_lm[:, 2] = mu_lm[:, 2] * -1 mu_lm = torch.stack(batch_size * [mu_lm.to(device=device)]) shape = mu_lm lm_eigenvec = torch.from_numpy( loader.lm_eigenvec).float().to(device=device) sigma = torch.from_numpy(loader.sigma).float().detach().to(device=device) sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) lm_eigenvec = torch.stack(M * [lm_eigenvec]) # main training loop best = 10000 for epoch in itertools.count(): for j, batch in enumerate(loader): # get the input and gt values x_cam_gt = batch['x_cam_gt'].to(device=device) shape_gt = batch['x_w_gt'].to(device=device) fgt = batch['f_gt'].to(device=device) x_img = batch['x_img'].to(device=device) #beta_gt = batch['beta_gt'].to(device=device) #x_img_norm = batch['x_img_norm'] #x_img_gt = batch['x_img_gt'].to(device=device).permute(0,2,1,3) x = x_img.reshape(M, N, 2).permute(0, 2, 1) batch_size = fgt.shape[0] #x_cam_pt = x_cam_gt.permute(0,1,3,2).reshape(batch_size,6800,3) #x = x_img.permute(0,2,1).reshape(batch_size,2,M,N) #ptsI = x_img_one.reshape(batch_size,M,N,3).permute(0,1,3,2)[:,:,:2,:] # calibration f = torch.squeeze(calib_net(x) + 300) K = torch.zeros((M, 3, 3)).float().to(device=device) K[:, 0, 0] = f K[:, 1, 1] = f K[:, 2, 2] = 1 # sfm betas = sfm_net(x) betas = betas.unsqueeze(-1) shape = mu_lm + torch.bmm(lm_eigenvec, betas).squeeze().view( M, N, 3) shape = shape - shape.mean(1).unsqueeze(1) opt1.zero_grad() opt2.zero_grad() f_error = torch.mean(torch.abs(f - fgt)) #error2d = torch.mean(torch.abs(pred - x_img_gt)) error3d = torch.mean(torch.norm(shape - shape_gt, dim=2)) error = f_error + error3d error.backward() opt1.step() opt2.step() print( f"iter: {j} | best: {best:.2f} | f_error: {f_error.item():.3f} | error3d: {error3d.item():.3f} " ) if j == 1000: break # save model and increment weight decay torch.save(sfm_net.state_dict(), os.path.join('model', 'sfm_model.pt')) torch.save(calib_net.state_dict(), os.path.join('model', 'calib_model.pt')) ferror = test(modelin='model.pt', outfile=args.out, optimize=False) if ferror < best: best = ferror print("saving!") torch.save(sfm_net.state_dict(), os.path.join('model', 'sfm_' + modelout)) torch.save(calib_net.state_dict(), os.path.join('model', 'calib_' + modelout)) sfm_net.train() calib_net.train()
def test(modelin=args.model, outfile=args.out, optimize=args.opt): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1) sfm_net = PointNet(n=199) if modelin != "": calib_path = os.path.join('model', 'calib_' + modelin) sfm_path = os.path.join('model', 'sfm_' + modelin) calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.eval() sfm_net.eval() 68 // 10 Mvals = [i for i in range(1, 100)] Nvals = [i for i in range(3, 68)] f_vals = [i * 200 for i in range(2, 7)] fpred_mean = np.zeros((100, 65, 5, 5)) fpred_med = np.zeros((100, 65, 5, 5)) fpred = np.zeros((100, 65, 5, 5)) factual = np.zeros((100, 65, 5, 5)) depth_error = np.zeros((100, 65, 5, 5)) for i, viewcount in enumerate(Mvals): for j, ptcount in enumerate(Nvals): for l, ftest in enumerate(f_vals): data3dmm = dataloader.MNLoader(M=viewcount, N=ptcount, f=ftest, seed=0) M = data3dmm.M N = data3dmm.N # mean shape and eigenvectors for 3dmm mu_s = torch.from_numpy(data3dmm.mu_s).float().detach() mu_s[:, 2] = mu_s[:, 2] * -1 lm_eigenvec = torch.from_numpy( data3dmm.lm_eigenvec).float().detach() sigma = torch.from_numpy(data3dmm.sigma).float().detach() sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) mu_s = torch.from_numpy(data3dmm.mu_s).float().detach() mu_s[:, 2] = mu_s[:, 2] * -1 lm_eigenvec = torch.from_numpy( data3dmm.lm_eigenvec).float().detach() sigma = torch.from_numpy(data3dmm.sigma).float().detach() sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) for k in range(5): data = data3dmm[k] # load the data x_cam_gt = data['x_cam_gt'] shape_gt = data['x_w_gt'] fgt = data['f_gt'] x_img = data['x_img'] x_img_gt = data['x_img_gt'] ptsI = x_img.reshape((M, N, 2)).permute(0, 2, 1) x = x_img.unsqueeze(0).permute(0, 2, 1) # run the model f = torch.squeeze(calib_net(ptsI) + 300) betas = sfm_net(ptsI) betas = betas.unsqueeze(-1) eigenvec = torch.stack(M * [lm_eigenvec]) shape = torch.stack(M * [mu_s]) + torch.bmm( eigenvec, betas).squeeze().view(M, N, 3) shape = shape - shape.mean(1).unsqueeze(1) shape = shape.mean(0) # get motion measurement guess K = torch.zeros((M, 3, 3)).float() K[:, 0, 0] = f K[:, 1, 1] = f K[:, 2, 2] = 1 km, c_w, scaled_betas, alphas = util.EPnP_single( ptsI, shape, K) _, R, T, mask = util.optimizeGN(km, c_w, scaled_betas, alphas, shape, ptsI) error_time = util.getTimeConsistency(shape, R, T) if error_time > 20: mode = 'walk' else: mode = 'still' # apply dual optimization if optimize: calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) shape, K, R, T = dualoptimization(x, calib_net, sfm_net, lm_eigenvec, betas, mu_s, shape_gt=shape_gt, fgt=fgt, M=M, N=N, mode=mode) f = K[0, 0].detach() # get motion measurement guess fmu = f.mean() fmed = f.flatten().median() K = torch.zeros(M, 3, 3).float() K[:, 0, 0] = fmu K[:, 1, 1] = fmu K[:, 2, 2] = 1 K, _ = torch.median(K, dim=0) km, c_w, scaled_betas, alphas = util.EPnP(ptsI, shape, K) Xc, R, T, mask = util.optimizeGN(km, c_w, scaled_betas, alphas, shape, ptsI) # get errors rel_errors = util.getRelReprojError3(x_cam_gt, shape, R, T) fpred_mean[i, j, l, k] = fmu.detach().cpu().item() fpred_med[i, j, l, k] = fmed.detach().cpu().item() factual[i, j, l, k] = fgt.detach().cpu().item() depth_error[i, j, l, k] = rel_errors.cpu().mean().item() print( f"M: {viewcount} | N: {ptcount} | f/fgt: {fpred_mean[i,j,l,k]:.2f}/{factual[i,j,l,k]}" ) ferror_mu = np.mean( np.abs(fpred_mean[i, j, l] - factual[i, j, l]) / factual[i, j, l]) ferror_med = np.mean( np.abs(fpred_med[i, j, l] - factual[i, j, l]) / factual[i, j, l]) derror = np.mean(depth_error[i, j, l]) f = np.mean(fpred_mean[i, j, l]) print( f"M: {viewcount} | N: {ptcount} | fgt: {ftest:.2f} | ferror_mu: {ferror_mu:.2f} | ferror_med: {ferror_med:.2f} | derror: {derror:.2f}" ) matdata = {} matdata['fpred_mu'] = fpred_mean matdata['fpred_med'] = fpred_med matdata['fgt'] = factual matdata['derror'] = depth_error scipy.io.savemat(outfile, matdata) print(f"saved output to {outfile}")
def testReal(modelin=args.model, outfile=args.out, optimize=args.opt, db=args.db): # define model, dataloader, 3dmm eigenvectors, optimization method calib_net = PointNet(n=1) sfm_net = PointNet(n=199) if modelin != "": calib_path = os.path.join('model', 'calib_' + modelin) sfm_path = os.path.join('model', 'sfm_' + modelin) calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) calib_net.eval() sfm_net.eval() # mean shape and eigenvectors for 3dmm data3dmm = dataloader.SyntheticLoader() mu_lm = torch.from_numpy(data3dmm.mu_lm).float().detach() mu_lm[:, 2] = mu_lm[:, 2] * -1 lm_eigenvec = torch.from_numpy(data3dmm.lm_eigenvec).float().detach() sigma = torch.from_numpy(data3dmm.sigma).float().detach() sigma = torch.diag(sigma.squeeze()) lm_eigenvec = torch.mm(lm_eigenvec, sigma) # define loader loader = getLoader(db) out_fpred = [] out_fgt = [] out_dpred = [] out_dgt = [] shape_pred = [] error_2d = [] error_relf = [] error_rel3d = [] for sub in range(len(loader)): batch = loader[sub] x_cam_gt = batch['x_cam_gt'] fgt = batch['f_gt'] x_img = batch['x_img'] x_img_gt = batch['x_img_gt'] M = x_img_gt.shape[0] N = x_img_gt.shape[-1] ptsI = x_img.reshape((M, N, 2)).permute(0, 2, 1) x = x_img.unsqueeze(0).permute(0, 2, 1) # run the model f = torch.squeeze(calib_net(ptsI) + 300) betas = sfm_net(ptsI) betas = betas.unsqueeze(-1) eigenvec = torch.stack(M * [lm_eigenvec]) shape = torch.stack(M * [mu_lm]) + torch.bmm( eigenvec, betas).squeeze().view(M, N, 3) shape = shape - shape.mean(1).unsqueeze(1) shape = shape.mean(0) # get motion measurement guess K = torch.zeros((M, 3, 3)).float() K[:, 0, 0] = f K[:, 1, 1] = f K[:, 2, 2] = 1 km, c_w, scaled_betas, alphas = util.EPnP_single(ptsI, shape, K) _, R, T, mask = util.optimizeGN(km, c_w, scaled_betas, alphas, shape, ptsI) error_time = util.getTimeConsistency(shape, R, T) if error_time > 20: mode = 'walk' else: mode = 'still' print(mode, error_time) # additional optimization on initial solution shape_gt = batch['x_w_gt'] if db == 'biwi' else None if optimize: calib_net.load_state_dict(torch.load(calib_path)) sfm_net.load_state_dict(torch.load(sfm_path)) print(mode) if db == 'biwi': shape, K, R, T, iter = dualoptimization(ptsI, calib_net, sfm_net, shape_gt=shape_gt, fgt=fgt, db='biwi', mode=mode) else: shape, K, R, T, iter = dualoptimization(ptsI, calib_net, sfm_net, fgt=fgt, mode=mode) f = K[:, 0, 0].detach() # get pose with single intrinsic fmu = f.mean() fmed = f.flatten().median() K = torch.zeros(M, 3, 3).float() K[:, 0, 0] = fmu K[:, 1, 1] = fmu K[:, 2, 2] = 1 km, c_w, scaled_betas, alphas = util.EPnP_single(ptsI, shape, K) Xc, R, T, mask = util.optimizeGN(km, c_w, scaled_betas, alphas, shape, ptsI) # get errors #reproj_errors2 = util.getError(ptsI,shape,R,T,K) reproj_errors2 = util.getReprojError2(ptsI, shape, R, T, K.mean(0), show=False, loss='l2') rel_errors = util.getRelReprojError3(x_cam_gt, shape, R, T) d = torch.norm(T, dim=1) dgt = torch.norm(torch.mean(x_cam_gt, dim=2), dim=1) reproj_error = reproj_errors2.mean() rel_error = rel_errors.mean() f_error = torch.mean(torch.abs(fgt - fmu) / fgt) # save final prediction out_fpred.append(f.detach().cpu().numpy()) out_fgt.append(fgt.numpy()) out_dpred.append(d.detach().cpu().numpy()) out_dgt.append(dgt.cpu().numpy()) f_x = torch.mean(fmu.detach()).cpu().item() shape_pred.append(shape.detach().cpu().numpy()) error_2d.append(reproj_error.cpu().data.item()) error_rel3d.append(rel_error.cpu().data.item()) error_relf.append(f_error.cpu().data.item()) print( f" f/fgt: {f_x:.3f}/{fgt.item():.3f} | f_error_rel: {f_error.item():.4f} | rel rmse: {rel_error.item():.4f} | 2d error: {reproj_error.item():.4f}" ) #end for # prepare output file out_shape = np.stack(shape_pred) out_fpred = np.array(out_fpred, dtype=np.object) out_fgt = np.array(out_fgt, dtype=np.object).T matdata = {} matdata['fpred'] = out_fpred matdata['fgt'] = out_fgt matdata['dpred'] = out_dpred matdata['dgt'] = out_dgt matdata['shape'] = np.stack(out_shape) matdata['error_2d'] = np.array(error_2d) matdata['error_rel3d'] = np.array(error_rel3d) matdata['error_relf'] = np.array(error_relf) scipy.io.savemat(outfile, matdata) print(f"MEAN seterror_2d: {np.mean(error_2d)}") print(f"MEAN seterror_rel3d: {np.mean(error_rel3d)}") print(f"MEAN seterror_relf: {np.mean(error_relf)}")