def validation(model, val_loader,p3d_mean, p3d_std):
joint_error = []
PA_MPJPE = []
with torch.no_grad():
model.eval()
for _, (_, pose3d, _, pose2d, _) in enumerate(val_loader):
p2d = pose2d.to(DEVICE).float()
p3d_pred = model(p2d)
# unnormalize groundtruth pose3d, after unnorm, pelvis = [0,0,0]
p3d_gt_unnorm = pose3d.cpu().detach().numpy().reshape([BATCHSIZE, -1, 3])
p3d_pelvis = p3d_gt_unnorm[:, 0, :]
p3d_pelvis = np.expand_dims(p3d_pelvis, axis=1)
p3d_pelvis = np.repeat(p3d_pelvis, 17, axis=1)
p3d_gt_unnorm = p3d_gt_unnorm - p3d_pelvis
# unnormalize predicted pose3d
p3d_pred_np = p3d_pred.cpu().detach().numpy().reshape([BATCHSIZE, -1, 3])
p3d_pred_unnorm = util.unnormalize(pose3d_norm=p3d_pred_np,
mean=p3d_mean,
std=p3d_std,
num_joints=17)
MPJPE = util.get_error(pose3d_pred=p3d_pred_unnorm, pose3d_gt=p3d_gt_unnorm)
joint_error.append(MPJPE[0])
PA_MPJPE.append(MPJPE[1])
joint_error_mean = np.array(joint_error).mean()
PA_MPJPE_mean = np.array(PA_MPJPE).mean()
return joint_error_mean, PA_MPJPE_mean
def evaluate(model, dataset, device, filename):
image, mask, gt = zip(*[dataset[i] for i in range(8)])
image = torch.stack(image)
mask = torch.stack(mask)
gt = torch.stack(gt)
with torch.no_grad():
output, _ = model(image.to(device), mask.to(device))
output = output.to(torch.device('cpu'))
output_comp = mask * image + (1 - mask) * output
grid = make_grid(
torch.cat((unnormalize(image), mask, unnormalize(output),
unnormalize(output_comp), unnormalize(gt)),
dim=0))
save_image(grid, filename)
def test_hrnet_semgcn(hrnet, semgcn, argmax, test_loader, p3d_mean, p3d_std):
hrnet.eval()
semgcn.eval()
for idx, img in enumerate(test_loader):
img = img.to(DEVICE)
# forward
pred_heatmap = hrnet(img)
pred_p2d = argmax(pred_heatmap).view([-1, 17, 2])
pred_p3d = semgcn(pred_p2d).view([BATCHSIZE_TEST,
-1]).cpu().detach().numpy().reshape(
[BATCHSIZE_TEST, -1, 3])
pred_p3d_unnorm = util.unnormalize(pose3d_norm=pred_p3d,
mean=p3d_mean,
std=p3d_std,
num_joints=17)
pred_p3d_unnorm = pred_p3d_unnorm.reshape([-1, 51])
# record
if not idx:
predictions = pred_p3d_unnorm
else:
predictions = np.concatenate([predictions, pred_p3d_unnorm],
axis=0)
return predictions
def train(model, train_loader, val_loader, optimizer, p3d_mean, p3d_std,
traincounter):
model.train()
val_joint_error_mean = 1000
val_PA_MPJPE_mean = 1000
loader = tqdm(train_loader)
for _, (_, pose3d, pose3d_norm, pose2d, _) in enumerate(loader):
p3d_gt_norm, p2d = pose3d_norm.to(DEVICE), pose2d.to(
DEVICE).float().view([-1, 17, 2])
# forward
p3d_pred = model(p2d).view([-1, 51])
loss = F.mse_loss(p3d_gt_norm, p3d_pred, reduction='mean')
# backward
optimizer.zero_grad()
loss.backward()
optimizer.step()
# unnormalize groundtruth pose3d
p3d_gt_unnorm = pose3d.cpu().detach().numpy().reshape(
[BATCHSIZE, -1, 3])
p3d_pelvis = p3d_gt_unnorm[:, 0, :]
p3d_pelvis = np.expand_dims(p3d_pelvis, axis=1)
p3d_pelvis = np.repeat(p3d_pelvis, 17, axis=1)
p3d_gt_unnorm = p3d_gt_unnorm - p3d_pelvis
# unnormalize predicted pose3d
p3d_pred_np = p3d_pred.cpu().detach().numpy().reshape(
[BATCHSIZE, -1, 3])
p3d_pred_unnorm = util.unnormalize(pose3d_norm=p3d_pred_np,
mean=p3d_mean,
std=p3d_std,
num_joints=17)
traincounter += 1
if not traincounter % 1000:
torch.save(model.state_dict(), MODEL_PATH + SEMGCN)
if not traincounter % 5000:
val_joint_error_mean, val_PA_MPJPE_mean = validation(
model=model,
val_loader=val_loader,
p3d_mean=p3d_mean,
p3d_std=p3d_std)
print('val_joint_error_mean: {}, val_PA_MPJPE_mean: {}'.format(
val_joint_error_mean, val_PA_MPJPE_mean))
MPJPE = util.get_error(pose3d_pred=p3d_pred_unnorm,
pose3d_gt=p3d_gt_unnorm)
loader.set_description("joint error:{:.4f}, MPJPE: {:.4f}".format(
MPJPE[0], MPJPE[1]))
loader.refresh()
return traincounter
def log_image(self, img, step):
"""Log image into tensorboard """
# get first image
img = img[0]
# unnormalize first image
img = util.unnormalize(img)
img[img > 1] = 1
img[img < 0] = 0
# log image
self.writer.add_image('images', img, step)
def Result_Plot(test_loader, model, device, img_mean, img_std):
#fig, ax = plt.subplots(nrows=10, ncols=5, figsize=(30,30))
fig = plt.figure(figsize=(30, 30))
fig.subplots_adjust(hspace=0.1, wspace=0.000001)
# Get the first batch
data = next(iter(test_loader))
input_data = torch.cat([data[0], data[1]], dim=1)
input_data = input_data.to(device)
mask_label = data[2].to(device)
depth_label = data[3].to(device)
# Predict
pred = model(input_data)
ctr = 1
#fig.suptitle("Result os the Trained Model on Test Images", fontweight="bold", fontsize=16)
for i in range(5):
plt.axis('off')
fg_bg = unnormalize(data[1][i], img_mean, img_std)
plt.subplot(10, 5, ctr)
plt.imshow(fg_bg)
if ctr < 6:
plt.title("FG_BG IMAGE", fontweight="bold", fontsize=20)
plt.axis('off')
ctr += 1
label_mask = unnormalize_1D(data[2][i])
plt.subplot(10, 5, ctr)
plt.imshow(label_mask, cmap='gray')
if ctr < 6:
plt.title("Ground truth: MASK", fontweight="bold", fontsize=20)
plt.axis('off')
ctr += 1
label_depth = unnormalize_1D(data[3][i])
plt.subplot(10, 5, ctr)
plt.imshow(label_depth, cmap='gray')
if ctr < 6:
plt.title("Ground truth: DEPTH", fontweight="bold", fontsize=20)
plt.axis('off')
ctr += 1
pred_mask = unnormalize_1D(pred[0][i])
plt.subplot(10, 5, ctr)
plt.imshow(pred_mask, cmap='gray')
if ctr < 6:
plt.title("Predicted MASK", fontweight="bold", fontsize=20)
plt.axis('off')
ctr += 1
pred_depth = unnormalize_1D(pred[1][i])
plt.subplot(10, 5, ctr)
plt.imshow(pred_depth, cmap='gray')
if ctr < 6:
plt.title("Predicted DEPTH", fontweight="bold", fontsize=20)
plt.axis('off')
ctr += 1
#plt.tight_layout()
#plt.subplots_adjust(wspace=0, hspace=0)
plt.savefig('/content/RESULT_fig.jpg')
plt.show()