def load_trainer_from_model(OUT_DIR, device, model_name, z_dim):
'''
Create Network, model saver and trainer.
:param OUT_DIR: Where the model is stored or where it should be stored if there is none
:param device: torch device
:param model_name: the namen under which the model gets saved
:param z_dim: dimesnianality of latent space
:return: model loader/saver, current epoch, iteration, trainer for the model
'''
occ_net = OccupancyNetwork(z_dim=z_dim, device=device)
glorot_weight_zero_bias(occ_net)
optimizer = optim.Adam(occ_net.parameters(), lr=1e-4)
# Restore the model
checkpoint_io = CheckpointIO(OUT_DIR, model=occ_net, optimizer=optimizer)
try:
load_dict = checkpoint_io.load(model_name)
except FileExistsError:
load_dict = dict()
epoch_it = load_dict.get('epoch_it', -1)
it = load_dict.get('it', -1)
# metric_val_best = load_dict.get(
# 'loss_val_best', -model_selection_sign * np.inf)
# TODO: Validation scores for the model
# Write to tensorboard
trainer = Trainer(occ_net, optimizer, device=device)
return checkpoint_io, epoch_it, it, trainer
def main(cfg):
if cfg['mode'] == 'train':
train_dataset = get_dataset(mode=cfg['mode'], cfg=cfg)
val_dataset = get_dataset(mode='val', cfg=cfg)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg['train']['batch_size'],
num_workers=8,
shuffle=True,
collate_fn=collate_remove_none,
worker_init_fn=worker_init_fn)
val_loader = torch.utils.data.DataLoader(
val_dataset,
batch_size=cfg['val']['batch_size'],
num_workers=8,
shuffle=False,
collate_fn=collate_remove_none,
worker_init_fn=worker_init_fn)
model = get_network(cfg, device='cuda:0', dataset=train_dataset)
else:
test_dataset = get_dataset(mode=cfg['mode'], cfg=cfg, return_idx=True)
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
num_workers=4,
shuffle=False)
model = get_network(cfg, device='cuda:0', dataset=test_dataset)
if cfg['mode'] == 'train':
optimizer = optim.Adam(model.parameters(), lr=1e-4)
else:
optimizer = None
if cfg['mode'] == 'train':
checkpoint = CheckpointIO(cfg['out']['checkpoint_dir'],
model=model,
optimizer=optimizer)
load_dict = checkpoint.load(cfg['train']['pretrained'])
train(train_loader, val_loader, model, optimizer, checkpoint, cfg)
else:
checkpoint = CheckpointIO(cfg['out']['checkpoint_dir'], model=model)
load_dict = checkpoint.load(cfg['test']['pretrained'])
test(test_loader, test_dataset, model, cfg)
def live_application(arg):
model = HandNet()
model = model.to(device)
checkpoint_io = CheckpointIO('.', model=model)
load_dict = checkpoint_io.load('checkpoints/model.pt')
model.eval()
dd = pickle.load(open("MANO_RIGHT.pkl", 'rb'), encoding='latin1')
face = np.array(dd['f'])
renderer = utils.MeshRenderer(face, img_size=256)
cx = arg.cx
cy = arg.cy
fx = arg.fx
fy = arg.fy
gr = jit(grad(residuals))
lr = 0.03
opt_init, opt_update, get_params = optimizers.adam(lr, b1=0.5, b2=0.5)
opt_init = jit(opt_init)
opt_update = jit(opt_update)
get_params = jit(get_params)
i = 0
img_list = glob.glob("./demo/*")
with torch.no_grad():
for img_path in img_list:
i = i + 1
img = np.array(Image.open(img_path))
if img is None:
continue
if img.shape[0] > img.shape[1]:
margin = int((img.shape[0] - img.shape[1]) / 2)
img = img[margin:-margin]
cy = cy - margin
width = img.shape[1]
elif img.shape[0] < img.shape[1]:
margin = int((img.shape[1] - img.shape[0]) / 2)
img = img[:, margin:-margin]
cx = cx - margin
width = img.shape[0]
img = cv2.resize(img, (256, 256), cv2.INTER_LINEAR)
frame = img.copy()
cx = (cx * 256) / width
cy = (cy * 256) / width
fx = (fx * 256) / width
fy = (fy * 256) / width
intr = torch.from_numpy(
np.array([
[fx, 0.0, cx],
[0.0, fy, cy],
[0.0, 0.0, 1.0],
],
dtype=np.float32)).unsqueeze(0).to(device)
_intr = intr.cpu().numpy()
camparam = np.zeros((1, 21, 4))
camparam[:, :, 0] = _intr[:, 0, 0]
camparam[:, :, 1] = _intr[:, 1, 1]
camparam[:, :, 2] = _intr[:, 0, 2]
camparam[:, :, 3] = _intr[:, 1, 2]
img = functional.to_tensor(img).float()
img = functional.normalize(img, [0.5, 0.5, 0.5], [1, 1, 1])
img = img.unsqueeze(0).to(device)
hm, so3, beta, joint_root, bone = model(img, intr)
kp2d = hm_to_kp2d(hm.detach().cpu().numpy()) * 4
so3 = so3[0].detach().cpu().float().numpy()
beta = beta[0].detach().cpu().float().numpy()
bone = bone[0].detach().cpu().numpy()
joint_root = joint_root[0].detach().cpu().numpy()
so3 = npj.array(so3)
beta = npj.array(beta)
bone = npj.array(bone)
joint_root = npj.array(joint_root)
kp2d = npj.array(kp2d)
so3_init = so3
beta_init = beta
joint_root = reinit_root(joint_root, kp2d, camparam)
joint = mano_de_j(so3, beta)
bone = reinit_scale(joint, kp2d, camparam, bone, joint_root)
params = {'so3': so3, 'beta': beta, 'bone': bone}
opt_state = opt_init(params)
n = 0
while n < 20:
n = n + 1
params = get_params(opt_state)
grads = gr(params, so3_init, beta_init, joint_root, kp2d,
camparam)
opt_state = opt_update(n, grads, opt_state)
params = get_params(opt_state)
v = mano_de(params, joint_root, bone)
frame1 = renderer(v, intr[0].cpu(), frame)
cv2.imwrite("./out/" + str(i) + "_input.png", np.flip(frame, -1))
cv2.imwrite("./out/" + str(i) + "_output.png", np.flip(frame1, -1))
def live_application():
model = HandNetInTheWild()
model = model.to(device)
checkpoint_io = CheckpointIO('.', model=model)
load_dict = checkpoint_io.load('checkpoints/model.pt')
model.eval()
dd = pickle.load(open("MANO_RIGHT.pkl", 'rb'), encoding='latin1')
face = np.array(dd['f'])
renderer = utils.OpendrRenderer(img_size=256)
gr = jit(grad(residuals))
lr = 0.03
opt_init, opt_update, get_params = optimizers.adam(lr, b1=0.5, b2=0.5)
opt_init = jit(opt_init)
opt_update = jit(opt_update)
get_params = jit(get_params)
i = 0
img_list = glob.glob("./demo/*")
with torch.no_grad():
for img_path in img_list:
i = i + 1
img = np.array(Image.open(img_path))
if img is None:
continue
if img.shape[0] > img.shape[1]:
margin = int((img.shape[0] - img.shape[1]) / 2)
img = img[margin:-margin]
elif img.shape[0] < img.shape[1]:
margin = int((img.shape[1] - img.shape[0]) / 2)
img = img[:, margin:-margin]
img = cv2.resize(img, (256, 256), cv2.INTER_LINEAR)
frame = img.copy()
img = functional.to_tensor(img).float()
img = functional.normalize(img, [0.5, 0.5, 0.5], [1, 1, 1])
img = img.unsqueeze(0).to(device)
hm, so3, beta, joint_root, bone = model(img)
kp2d = hm_to_kp2d(hm.detach().cpu().numpy()) * 4
t = kp2d[0, 9, :]
so3 = so3[0].detach().cpu().float().numpy()
beta = beta[0].detach().cpu().float().numpy()
bone = bone[0].detach().cpu().numpy()
joint_root = joint_root[0].detach().cpu().numpy()
so3 = npj.array(so3)
beta = npj.array(beta)
bone = npj.array(bone)
joint_root = npj.array(joint_root)
kp2d = npj.array(kp2d)
so3_init = so3
beta_init = beta
joint = mano_de_j(so3, beta)
s = reinit_scale(joint, kp2d, t)
params = {'so3': so3, 'beta': beta}
so3 = params['so3']
beta = params['beta']
opt_state = opt_init(params)
n = 0
while n < 20:
n = n + 1
params = get_params(opt_state)
grads = gr(params, so3_init, beta_init, kp2d, s, t)
opt_state = opt_update(n, grads, opt_state)
joint = mano_de_j(so3, beta)
s = reinit_scale(joint, kp2d, t)
params = get_params(opt_state)
so3 = params['so3']
beta = params['beta']
v = mano_de(so3, beta)
v = v * s + np.array([t[0], t[1], 0])
frame1 = renderer.render(v.copy(), face, frame)
cv2.imwrite("./out/" + str(i) + "_input.png", np.flip(frame, -1))
cv2.imwrite("./out/" + str(i) + "_output.png", np.flip(frame1, -1))
def live_application(capture,arg):
window_size = 768
pygame.init()
display = pygame.display.set_mode((window_size, window_size))
pygame.display.set_caption('Real Time Hand Recon')
model = HandNet()
model = model.to(device)
checkpoint_io = CheckpointIO('.', model=model)
load_dict = checkpoint_io.load('checkpoints/model.pt')
model.eval()
dd = pickle.load(open("MANO_RIGHT.pkl", 'rb'), encoding='latin1')
face = np.array(dd['f'])
renderer = utils.MeshRenderer(face, img_size=256)
cx = arg.cx
cy = arg.cy
fx = arg.fx
fy = arg.fy
while True:
img = capture.read()
if img is None:
continue
if img.shape[0] > img.shape[1]:
margin = int((img.shape[0] - img.shape[1]) / 2)
cy = cy - margin
width = img.shape[1]
elif img.shape[0] < img.shape[1]:
margin = int((img.shape[1] - img.shape[0]) / 2)
cx = cx - margin
width = img.shape[0]
cx = (cx * 256)/width
cy = (cy * 256)/width
fx = (fx * 256)/width
fy = (fy * 256)/width
break
intr = torch.from_numpy(np.array([
[fx, 0.0, cx],
[0.0, fy, cy],
[0.0, 0.0, 1.0],
], dtype=np.float32)).unsqueeze(0).to(device)
_intr = intr.cpu().numpy()
camparam = np.zeros((1, 21, 4))
camparam[:, :, 0] = _intr[:, 0, 0]
camparam[:, :, 1] = _intr[:, 1, 1]
camparam[:, :, 2] = _intr[:, 0, 2]
camparam[:, :, 3] = _intr[:, 1, 2]
gr = jit(grad(residuals))
lr = 0.03
opt_init, opt_update, get_params = optimizers.adam(lr, b1=0.5, b2=0.5)
opt_init = jit(opt_init)
opt_update = jit(opt_update)
get_params = jit(get_params)
i = 0
with torch.no_grad():
while True:
i = i + 1
img = capture.read()
if img is None:
continue
if img.shape[0] > img.shape[1]:
margin = int((img.shape[0] - img.shape[1]) / 2)
img = img[margin:-margin]
elif img.shape[0] < img.shape[1]:
margin = int((img.shape[1] - img.shape[0]) / 2)
img = img[:, margin:-margin]
img = cv2.resize(img, (256, 256),cv2.INTER_LINEAR)
frame = img.copy()
img = functional.to_tensor(img).float()
img = functional.normalize(img, [0.5, 0.5, 0.5], [1, 1, 1])
img = img.unsqueeze(0).to(device)
hm, so3, beta, joint_root, bone = model(img,intr)
kp2d = hm_to_kp2d(hm.detach().cpu().numpy())*4
so3 = so3[0].detach().cpu().float().numpy()
beta = beta[0].detach().cpu().float().numpy()
bone = bone[0].detach().cpu().numpy()
joint_root = joint_root[0].detach().cpu().numpy()
so3 = npj.array(so3)
beta = npj.array(beta)
bone = npj.array(bone)
joint_root = npj.array(joint_root)
kp2d = npj.array(kp2d)
so3_init = so3
beta_init = beta
joint_root = reinit_root(joint_root,kp2d, camparam)
joint = mano_de_j(so3, beta)
bone = reinit_scale(joint,kp2d,camparam,bone,joint_root)
params = {'so3':so3, 'beta':beta, 'bone':bone}
opt_state = opt_init(params)
n = 0
while n < 20:
n = n + 1
params = get_params(opt_state)
grads = gr(params,so3_init,beta_init,joint_root,kp2d,camparam)
opt_state = opt_update(n, grads, opt_state)
params = get_params(opt_state)
v = mano_de(params,joint_root,bone)
frame = renderer(v,intr[0].cpu(),frame)
display.blit(
pygame.surfarray.make_surface(np.transpose(cv2.resize(np.flip(frame,1), (window_size, window_size),cv2.INTER_LINEAR), (1, 0, 2))),(0, 0))
pygame.display.update()
model_name = 'model' + '_z_dim_' + str(z_dim) + '.pt'
DATASET_PATH = os.path.join(current_dir, data_path, voxel_model, '')
print(DATASET_PATH)
OUT_DIR = os.path.join(current_dir, out_path, voxel_model, '')
GEN_DIR = os.path.join(OUT_DIR, gen, '')
print(OUT_DIR)
if not os.path.exists(GEN_DIR):
os.makedirs(GEN_DIR)
# Create torch device for GPU computing
is_cuda = (torch.cuda.is_available())
device = torch.device("cuda" if is_cuda else "cpu")
# Create/Load model
occ_net = OccupancyNetwork(device=device, z_dim=z_dim)
checkpoint_io = CheckpointIO(OUT_DIR, model=occ_net)
iteration = 0
try:
load_dict = checkpoint_io.load(model_name)
iteration = load_dict
except FileExistsError:
print("No model found!")
load_dict = dict()
epoch_it = load_dict.get('epoch_it', -1)
it = load_dict.get('it', -1)
test_dataset = get_dataset("test", dataset_path=DATASET_PATH)
# Create the dataloader
test_loader = torch.utils.data.DataLoader(test_dataset,
batch_size=1,
def live_application(capture):
window_size = 768
pygame.init()
display = pygame.display.set_mode((window_size, window_size))
pygame.display.set_caption('Real Time Hand Recon')
model = HandNetInTheWild()
model = model.to(device)
checkpoint_io = CheckpointIO('.', model=model)
load_dict = checkpoint_io.load('checkpoints/model.pt')
model.eval()
dd = pickle.load(open("MANO_RIGHT.pkl", 'rb'), encoding='latin1')
face = np.array(dd['f'])
renderer = utils.OpendrRenderer(img_size=256)
gr = jit(grad(residuals))
lr = 0.03
opt_init, opt_update, get_params = optimizers.adam(lr, b1=0.5, b2=0.5)
opt_init = jit(opt_init)
opt_update = jit(opt_update)
get_params = jit(get_params)
with torch.no_grad():
while True:
img = capture.read()
if img is None:
continue
if img.shape[0] > img.shape[1]:
margin = int((img.shape[0] - img.shape[1]) / 2)
img = img[margin:-margin]
elif img.shape[0] < img.shape[1]:
margin = int((img.shape[1] - img.shape[0]) / 2)
img = img[:, margin:-margin]
img = cv2.resize(img, (256, 256), cv2.INTER_LINEAR)
frame = img.copy()
img = functional.to_tensor(img).float()
img = functional.normalize(img, [0.5, 0.5, 0.5], [1, 1, 1])
img = img.unsqueeze(0).to(device)
hm, so3, beta, joint_root, bone = model(img)
kp2d = hm_to_kp2d(hm.detach().cpu().numpy()) * 4
t = kp2d[0, 9, :]
so3 = so3[0].detach().cpu().float().numpy()
beta = beta[0].detach().cpu().float().numpy()
bone = bone[0].detach().cpu().numpy()
joint_root = joint_root[0].detach().cpu().numpy()
so3 = npj.array(so3)
beta = npj.array(beta)
bone = npj.array(bone)
joint_root = npj.array(joint_root)
kp2d = npj.array(kp2d)
so3_init = so3
beta_init = beta
joint = mano_de_j(so3, beta)
s = reinit_scale(joint, kp2d, t)
params = {'so3': so3, 'beta': beta}
so3 = params['so3']
beta = params['beta']
opt_state = opt_init(params)
n = 0
while n < 20:
n = n + 1
params = get_params(opt_state)
grads = gr(params, so3_init, beta_init, kp2d, s, t)
opt_state = opt_update(n, grads, opt_state)
joint = mano_de_j(so3, beta)
s = reinit_scale(joint, kp2d, t)
params = get_params(opt_state)
so3 = params['so3']
beta = params['beta']
v = mano_de(so3, beta)
v = v * s + np.array([t[0], t[1], 0])
frame = renderer.render(v.copy(), face, frame)
display.blit(
pygame.surfarray.make_surface(
np.transpose(
cv2.resize(np.flip(frame,
1), (window_size, window_size),
cv2.INTER_LINEAR), (1, 0, 2))), (0, 0))
pygame.display.update()
def live_application(capture, arg):
pygame.init()
display = pygame.display.set_mode((640, 480))
pygame.display.set_caption('Real Time Hand Recon')
dd = pickle.load(open("MANO_RIGHT.pkl", 'rb'), encoding='latin1')
face = np.array(dd['f'])
model = HandNet()
model = model.to(device)
checkpoint_io = CheckpointIO('.', model=model)
load_dict = checkpoint_io.load('checkpoints/model.pt')
model.eval()
renderer = utils.MeshRenderer(face, img_size=[640, 480])
o_intr = torch.from_numpy(
np.array([
[arg.fx, 0.0, arg.cx],
[0.0, arg.fy, arg.cy],
[0.0, 0.0, 1.0],
],
dtype=np.float32)).unsqueeze(0).numpy()
o_camparam = np.zeros((4))
o_camparam[0] = o_intr[0, 0, 0]
o_camparam[1] = o_intr[0, 1, 1]
o_camparam[2] = o_intr[0, 0, 2]
o_camparam[3] = o_intr[0, 1, 2]
gr = jit(grad(residuals))
lr = 0.03
opt_init, opt_update, get_params = optimizers.adam(lr, b1=0.5, b2=0.5)
opt_init = jit(opt_init)
opt_update = jit(opt_update)
get_params = jit(get_params)
x_reg = np.ones((10, )) * 240
y_reg = np.ones((10, )) * 240
s_reg = np.ones((10, )) * 240
weight = np.array([0, 0, 0, 0, 0, 0, 0, 0.1, 0.2, 0.7])
i = 0
x = 240
y = 320
scale = 256
with torch.no_grad():
while True:
i = i + 1
img = capture.read()
frame = img.copy()
if img is None:
continue
vmin = max(0, y - scale // 2)
vmin_p = max(scale // 2 - y, 0)
umin = max(0, x - scale // 2)
umin_p = max(scale // 2 - x, 0)
vmax = min(640, y + scale // 2)
vmax_p = max(scale // 2 + y - 640, 0)
umax = min(480, x + scale // 2)
umax_p = max(scale // 2 + x - 480, 0)
img = img[int(umin):int(umax), int(vmin):int(vmax), :]
img = cv2.copyMakeBorder(img,
int(umin_p),
int(umax_p),
int(vmin_p),
int(vmax_p),
cv2.BORDER_CONSTANT,
value=[255, 255, 255])
cx = arg.cx - y + scale // 2
cy = arg.cy - x + scale // 2
cx = (cx * 256) / scale
cy = (cy * 256) / scale
fx = (arg.fx * 256) / scale
fy = (arg.fy * 256) / scale
intr = torch.from_numpy(
np.array([
[fx, 0.0, cx],
[0.0, fy, cy],
[0.0, 0.0, 1.0],
],
dtype=np.float32)).unsqueeze(0).to(device)
_intr = intr.cpu().numpy()
camparam = np.zeros((1, 21, 4))
camparam[:, :, 0] = _intr[:, 0, 0]
camparam[:, :, 1] = _intr[:, 1, 1]
camparam[:, :, 2] = _intr[:, 0, 2]
camparam[:, :, 3] = _intr[:, 1, 2]
img = cv2.resize(img, (256, 256), cv2.INTER_LINEAR)
img = functional.to_tensor(img).float()
img = functional.normalize(img, [0.5, 0.5, 0.5], [1, 1, 1])
img = img.unsqueeze(0).to(device)
hm, so3, beta, joint_root, bone = model(img, intr)
kp2d = hm_to_kp2d(hm.detach().cpu().numpy()) * 4
so3 = so3[0].detach().cpu().float().numpy()
beta = beta[0].detach().cpu().float().numpy()
bone = bone[0].detach().cpu().numpy()
joint_root = joint_root[0].detach().cpu().numpy()
so3 = npj.array(so3)
beta = npj.array(beta)
bone = npj.array(bone)
joint_root = npj.array(joint_root)
kp2d = npj.array(kp2d)
so3_init = so3
beta_init = beta
joint_root = reinit_root(joint_root, kp2d, camparam)
joint = mano_de_j(so3, beta)
bone = reinit_scale(joint, kp2d, camparam, bone, joint_root)
params = {'so3': so3, 'beta': beta, 'bone': bone}
opt_state = opt_init(params)
n = 0
while n < 20:
n = n + 1
params = get_params(opt_state)
grads = gr(params, so3_init, beta_init, joint_root, kp2d,
camparam)
opt_state = opt_update(n, grads, opt_state)
params = get_params(opt_state)
v = mano_de(params, joint_root, bone)
kp2d = np.array(kp2d[0])
x = x + ((kp2d[9, 1] - 128) * scale) / 256
y = y + ((kp2d[9, 0] - 128) * scale) / 256
scale = max(
max(kp2d[:, 0].max() - kp2d[:, 0].min(),
kp2d[:, 1].max() - kp2d[:, 1].min()) * 2, 80)
x_reg[:9] = x_reg[1:]
x_reg[-1] = x
y_reg[:9] = y_reg[1:]
y_reg[-1] = y
s_reg[:9] = s_reg[1:]
s_reg[-1] = scale
x = (x_reg * weight).sum()
y = (y_reg * weight).sum()
scale = (s_reg * weight).sum()
frame = renderer(v, o_intr[0], frame)
frame = cv2.rectangle(frame, (int(vmin), int(umin)),
(int(vmax), int(umax)), (255, 255, 255),
thickness=5)
display.blit(
pygame.surfarray.make_surface(
np.transpose(np.flip(frame, 1), (1, 0, 2))), (0, 0))
pygame.display.update()