in bodymocap/train/eftFitter.py [0:0]
def run_eft_step_wHand(self, input_batch):
self.model.train()
if self.options.bExemplarMode:
self.exemplerTrainingMode()
# Get data from the batch
images = input_batch['img'] # input image
gt_keypoints_2d = input_batch['keypoints'].clone()# 2D keypoints #[N,49,3]
gt_pose = input_batch['pose'] # SMPL pose parameters #[N,72]
gt_betas = input_batch['betas'] # SMPL beta parameters #[N,10]
#3D joint scaling
JOINT_SCALING_3D = 1.0
gt_joints = input_batch['pose_3d']*JOINT_SCALING_3D # 3D pose #[N,24,4]
# has_smpl = input_batch['has_smpl'].byte() ==1 # flag that indicates whether SMPL parameters are valid
has_pose_3d = input_batch['has_pose_3d'].byte()==1 # flag that indicates whether 3D pose is valid
indices = input_batch['sample_index'] # index of example inside its dataset
batch_size = images.shape[0]
is_flipped = input_batch['is_flipped'] # flag that indicates whether image was flipped during data augmentation
rot_angle = input_batch['rot_angle'] # rotation angle used for data augmentation
dataset_name = input_batch['dataset_name'] # name of the dataset the image comes from
if self.options.bUseHand3D: #Hands
gt_lhand_3d = input_batch['lhand_3d']
gt_rhand_3d = input_batch['rhand_3d']
# Get GT vertices and model joints
# Note that gt_model_joints is different from gt_joints as it comes from SMPL
gt_out = self.smpl(betas=gt_betas, body_pose=gt_pose[:,3:], global_orient=gt_pose[:,:3])
gt_model_joints = gt_out.joints #[N, 49, 3] #Note this is different from gt_joints!
gt_vertices = gt_out.vertices
# Get current best fits from the dictionary
index_cpu = indices.cpu()
if self.options.bExemplar_dataLoaderStart>=0:
index_cpu +=self.options.bExemplar_dataLoaderStart #Bug fixed.
#Check existing SPIN fits
opt_pose, opt_betas, opt_validity = self.fits_dict[(dataset_name, index_cpu, rot_angle.cpu(), is_flipped.cpu())]
opt_pose = opt_pose.to(self.device)
opt_betas = opt_betas.to(self.device)
# De-normalize 2D keypoints from [-1,1] to pixel space
gt_keypoints_2d_orig = gt_keypoints_2d.clone()
gt_keypoints_2d_orig[:, :, :-1] = 0.5 * self.options.img_res * (gt_keypoints_2d_orig[:, :, :-1] + 1) #49: (25+24) x 3
# if gt_keypoints_2d[0,20,-1]>0: #Checking foot keypoint
# print(gt_keypoints_2d_orig[0,19:25])
# Feed images in the network to predict camera and SMPL parameters
pred_rotmat, pred_betas, pred_camera = self.model(images)
pred_output = self.smpl(betas=pred_betas, body_pose=pred_rotmat[:,1:], global_orient=pred_rotmat[:,0].unsqueeze(1), pose2rot=False)
pred_vertices = pred_output.vertices
pred_joints_3d = pred_output.joints
if self.options.bUseSMPLX:
pred_right_hand_joints_3d = pred_output.right_hand_joints
pred_left_hand_joints_3d = pred_output.left_hand_joints
# Convert Weak Perspective Camera [s, tx, ty] to camera translation [tx, ty, tz] in 3D given the bounding box size
# This camera translation can be used in a full perspective projection
pred_cam_t = torch.stack([pred_camera[:,1],
pred_camera[:,2],
2*self.focal_length/(self.options.img_res * pred_camera[:,0] +1e-9)],dim=-1)
camera_center = torch.zeros(batch_size, 2, device=self.device)
# weakProjection_gpu################
pred_keypoints_2d = weakProjection_gpu(pred_joints_3d, pred_camera[:,0], pred_camera[:,1:] ) #N, 49, 2
if True: #Ignore hips and hip centers, foot
LENGTH_THRESHOLD = 0.0089 #1/112.0 #at least it should be 5 pixel
#Disable parts
if True:
gt_keypoints_2d[:,2+25,2]=0
gt_keypoints_2d[:,3+25,2]=0
gt_keypoints_2d[:,14+25,2]=0
# #Compute angle knee to ankle orientation
gt_boneOri_leftLeg = gt_keypoints_2d[:,5+25,:2] - gt_keypoints_2d[:,4+25,:2] #Left lower leg orientation #(N,2)
gt_boneOri_leftLeg, leftLegLeng = normalize_2dvector(gt_boneOri_leftLeg)
if leftLegLeng>LENGTH_THRESHOLD:
leftLegValidity = gt_keypoints_2d[:,5+25, 2] * gt_keypoints_2d[:,4+25, 2]
pred_boneOri_leftLeg = pred_keypoints_2d[:,5+25,:2] - pred_keypoints_2d[:,4+25,:2]
pred_boneOri_leftLeg, _ = normalize_2dvector(pred_boneOri_leftLeg)
loss_legOri_left = torch.ones(1).to(self.device) - torch.dot(gt_boneOri_leftLeg.view(-1),pred_boneOri_leftLeg.view(-1))
else:
loss_legOri_left = torch.zeros(1).to(self.device)
leftLegValidity = torch.zeros(1).to(self.device)
gt_boneOri_rightLeg = gt_keypoints_2d[:,0+25,:2] - gt_keypoints_2d[:,1+25,:2] #Right lower leg orientation
gt_boneOri_rightLeg, rightLegLeng = normalize_2dvector(gt_boneOri_rightLeg)
if rightLegLeng>LENGTH_THRESHOLD:
rightLegValidity = gt_keypoints_2d[:,0+25, 2] * gt_keypoints_2d[:,1+25, 2]
pred_boneOri_rightLeg = pred_keypoints_2d[:,0+25,:2] - pred_keypoints_2d[:,1+25,:2]
pred_boneOri_rightLeg, _ = normalize_2dvector(pred_boneOri_rightLeg)
loss_legOri_right = torch.ones(1).to(self.device) - torch.dot(gt_boneOri_rightLeg.view(-1),pred_boneOri_rightLeg.view(-1))
else:
loss_legOri_right = torch.zeros(1).to(self.device)
rightLegValidity = torch.zeros(1).to(self.device)
# print("leftLegLeng: {}, rightLegLeng{}".format(leftLegLeng,rightLegLeng ))
loss_legOri = leftLegValidity* loss_legOri_left + rightLegValidity* loss_legOri_right
# loss_legOri = torch.zeros(1).to(self.device)
# if leftLegValidity.item():
# loss_legOri = loss_legOri + (pred_boneOri_leftLeg).mean()
# if rightLegValidity.item():
# loss_legOri = loss_legOri + self.criterion_regr(gt_boneOri_rightLeg, pred_boneOri_rightLeg)
# print(loss_legOri)
#Disable Foots
# gt_keypoints_2d[:,5+25,2]=0 #Left foot
# gt_keypoints_2d[:,0+25,2]=0 #Right foot
# Compute 2D reprojection loss for the keypoints
loss_keypoints_2d = self.keypoint_loss(pred_keypoints_2d, gt_keypoints_2d,
self.options.openpose_train_weight,
self.options.gt_train_weight)
# # Compute 3D keypoint loss
# loss_keypoints_3d = self.keypoint_3d_loss(pred_joints_3d, gt_joints, has_pose_3d)
loss_keypoints_3d = self.keypoint_3d_loss_panopticDB(pred_joints_3d, gt_joints, has_pose_3d)
# # loss_keypoints_3d = self.keypoint_3d_loss_modelSkel(pred_joints_3d, gt_model_joints[:,25:,:], has_pose_3d)
loss_regr_betas_noReject = torch.mean(pred_betas**2)
#Prevent bending knee?
# red_rotmat[0,6,:,:] -
loss = self.options.keypoint_loss_weight * loss_keypoints_2d + \
self.options.beta_loss_weight * loss_regr_betas_noReject + \
((torch.exp(-pred_camera[:,0]*10)) ** 2 ).mean()
if self.options.bExemplarWith3DSkel:
loss = loss + self.options.keypoint_loss_weight * loss_keypoints_3d
# loss = loss_keypoints_3d #TODO: DEBUGGIN
##### Compute 3D hand joint loss especially for panoptic stuido
if self.options.bUseHand3D:
loss_keypoints_3d_hand = self.keypoint_3d_hand_loss_panopticDB(pred_right_hand_joints_3d, pred_left_hand_joints_3d, gt_lhand_3d, gt_rhand_3d)
loss = loss + self.options.keypoint_loss_weight * loss_keypoints_3d_hand
if True: #Leg orientation loss
loss = loss + 0.005*loss_legOri
#Put zeor preference on knees
#Disabled. Not working
# if self.options.bUseKneePrior:
# eyemat = torch.eye(3,3).repeat(pred_rotmat.shape[0],1).view((-1,3,3)).to(self.device)
# kneePrior = ((pred_rotmat[:,5,:,:] - eyemat)**2).mean() + ((pred_rotmat[:,4,:,:] - eyemat)**2).mean()
# loss = loss + kneePrior*0.001
# print(loss_regr_betas)
loss *= 60
# print("loss2D: {}, loss3D: {}".format( self.options.keypoint_loss_weight * loss_keypoints_2d,self.options.keypoint_loss_weight * loss_keypoints_3d ) )
# Do backprop
self.optimizer.zero_grad()
loss.backward()
# g_timer.tic()
self.optimizer.step()
# g_timer.toc(bPrint =True)
# Pack output arguments for tensorboard logging
output = {'pred_vertices': 0, #pred_vertices.detach(),
'opt_vertices': 0,
'pred_cam_t': 0,#pred_cam_t.detach(),
'opt_cam_t': 0}
#Save result
output={}
output['pred_pose_rotmat'] = pred_rotmat.detach().cpu().numpy()
output['pred_shape'] = pred_betas.detach().cpu().numpy()
output['pred_camera'] = pred_camera.detach().cpu().numpy()
#If there exists SPIN fits, save that for comparison later
output['opt_pose'] = opt_pose.detach().cpu().numpy()
output['opt_beta'] = opt_betas.detach().cpu().numpy()
output['sampleIdx'] = input_batch['sample_index'].detach().cpu().numpy() #To use loader directly
output['imageName'] = input_batch['imgname']
output['scale'] = input_batch['scale'] .detach().cpu().numpy()
output['center'] = input_batch['center'].detach().cpu().numpy()
if 'annotId' in input_batch.keys():
output['annotId'] = input_batch['annotId'].detach().cpu().numpy()
if 'subjectId' in input_batch.keys():
output['subjectId'] = input_batch['subjectId'][0].item()
# print(output['subjectId'])
#To save new db file
output['keypoint2d'] = input_batch['keypoints_original'].detach().cpu().numpy()
output['keypoint2d_cropped'] = input_batch['keypoints'].detach().cpu().numpy()
#Save GT gt_lhand_3d, gt_rhand_3d
if self.options.bUseHand3D:
output['gt_lhand_3d'] = gt_lhand_3d.detach().cpu().numpy()
output['gt_rhand_3d'] = gt_rhand_3d.detach().cpu().numpy()
losses = {'loss': loss.detach().item(),
'loss_keypoints': loss_keypoints_2d.detach().item(),
'loss_keypoints_3d': loss_keypoints_3d.detach().item(),
# 'loss_regr_pose': loss_regr_pose.detach().item(),
'loss_regr_betas': loss_regr_betas_noReject.detach().item()}
# 'loss_shape': loss_shape.detach().item()}
if self.options.bDebug_visEFT:#g_debugVisualize: #Debug Visualize input
# print("Image Name: {}".format(output['imageName']))
for b in range(batch_size):
#denormalizeImg
curImgVis = images[b] #3,224,224
curImgVis = self.de_normalize_img(curImgVis).cpu().numpy()
curImgVis = np.transpose( curImgVis , (1,2,0) )*255.0
curImgVis =curImgVis[:,:,[2,1,0]]
#Denormalize image
curImgVis = np.ascontiguousarray(curImgVis, dtype=np.uint8)
originalImgVis = curImgVis.copy()
viewer2D.ImShow(curImgVis, name='rawIm')
gt_keypoints_2d_orig_vis = gt_keypoints_2d_orig.detach().cpu().numpy()
curImgVis = viewer2D.Vis_Skeleton_2D_SPIN49(gt_keypoints_2d_orig_vis[b,:,:2], gt_keypoints_2d_orig_vis[b,:,2], bVis= False, image=curImgVis)
# curImgVis = viewer2D.Vis_Skeleton_2D_Openpose18(gt_keypoints_2d_orig[b,:,:2].cpu().numpy(), gt_keypoints_2d_orig[b,:,2], bVis= False, image=curImgVis)
# Show projection of SMPL sksleton
if False:
pred_keypoints_2d_vis = pred_keypoints_2d[b,:,:2].detach().cpu().numpy()
pred_keypoints_2d_vis = 0.5 * self.options.img_res * (pred_keypoints_2d_vis + 1) #49: (25+24) x 3
if glViewer.g_bShowSkeleton:
curImgVis = viewer2D.Vis_Skeleton_2D_general(pred_keypoints_2d_vis, bVis= False, image=curImgVis)
viewer2D.ImShow(curImgVis, scale=2.0, waitTime=1)
#Get camera pred_params
pred_camera_vis = pred_camera.detach().cpu().numpy()
############### Visualize Mesh ###############
pred_vert_vis = pred_vertices[b].detach().cpu().numpy()
camParam_scale = pred_camera_vis[b,0]
camParam_trans = pred_camera_vis[b,1:]
pred_vert_vis = convert_smpl_to_bbox(pred_vert_vis, camParam_scale, camParam_trans)
pred_meshes = {'ver': pred_vert_vis, 'f': self.smpl.faces}
# glViewer.setMeshData([pred_meshes, opt_meshes], bComputeNormal= True)
glViewer.setMeshData([pred_meshes], bComputeNormal= True)
# glViewer.setMeshData([pred_meshes, opt_meshes], bComputeNormal= True)
# glViewer.setMeshData([opt_meshes], bComputeNormal= True)
# glViewer.SetMeshColor('red')
############### Visualize Skeletons ###############
#Vis pred-SMPL joint
pred_joints_vis = pred_joints_3d[b,:,:3].detach().cpu().numpy() #[N,49,3]
# pred_joints_pelvis = (pred_joints_vis[25+2,:] + pred_joints_vis[25+3,:]) / 2
pred_joints_nose = pred_joints_vis[25+19,:]# + pred_joints_vis[25+3,:]) / 2
pred_joints_vis = convert_smpl_to_bbox(pred_joints_vis, camParam_scale, camParam_trans)
pred_joints_vis = pred_joints_vis.ravel()[:,np.newaxis]
glViewer.setSkeleton( [pred_joints_vis], colorRGB=(0,0, 255), jointType='spin')
if self.options.bUseSMPLX:
for pred_rl in [pred_right_hand_joints_3d, pred_left_hand_joints_3d]:
pred_joints_vis = pred_rl[b,:,:3].detach().cpu().numpy() #[N,49,3]
pred_joints_vis = convert_smpl_to_bbox(pred_joints_vis, camParam_scale, camParam_trans)
pred_joints_vis = pred_joints_vis.ravel()[:,np.newaxis]
glViewer.addSkeleton( [pred_joints_vis], colorRGB=(0,0, 255), jointType='hand_panopticdb') #GT: blue
# #Vis GT joint (not model (SMPL) joint!!)
if has_pose_3d[b] and self.options.bExemplarWith3DSkel:
# gt_jointsVis = gt_model_joints[b,:,:3].cpu().numpy() #[N,49,3]
gt_jointsVis = gt_joints[b,:,:3].cpu().numpy() #[N,49,3]
# gt_pelvis = (gt_jointsVis[ 2,:] + gt_jointsVis[ 3,:]) / 2
gt_jointsVis = convert_smpl_to_bbox(gt_jointsVis, camParam_scale, camParam_trans)
gt_nose = gt_jointsVis[ 19,:]
gt_jointsVis = gt_jointsVis- gt_nose + pred_joints_nose
gt_jointsVis = gt_jointsVis.ravel()[:,np.newaxis]
# gt_jointsVis*=pred_camera_vis[b,0]
# gt_jointsVis[::3] += pred_camera_vis[b,1]
# gt_jointsVis[1::3] += pred_camera_vis[b,2]
# gt_jointsVis*=112
glViewer.addSkeleton( [gt_jointsVis], colorRGB=(255, 0, 0), jointType='spin')
#Vis hand
if self.options.bUseHand3D:
for gt_rl in [gt_lhand_3d, gt_rhand_3d]:
gt_3d = gt_rl[b,:,:3].cpu().numpy() #[N,49,3]
gt_3d = convert_smpl_to_bbox(gt_3d, camParam_scale, camParam_trans)
gt_3d = gt_3d - gt_nose + pred_joints_nose
gt_3d = gt_3d.ravel()[:,np.newaxis]
glViewer.addSkeleton( [gt_3d], colorRGB=(255,0,0), jointType='hand_panopticdb') #GT: red
# # glViewer.show()
glViewer.setBackgroundTexture(curImgVis) #Vis raw video as background
# glViewer.setBackgroundTexture(originalImgVis) #Vis raw video as background
glViewer.setWindowSize(curImgVis.shape[1]*3, curImgVis.shape[0]*3)
glViewer.SetOrthoCamera(True)
glViewer.show(1)
# glViewer.show_SMPL(bSaveToFile = True, bResetSaveImgCnt = False, mode = 'camera')
# glViewer.show_SMPL(bSaveToFile = True, bResetSaveImgCnt = False, mode = 'free')
# continue
return output, losses