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)}")
