Update sample/reconstruct_data_taco.py

sample/reconstruct_data_taco.py

@@ -51,15 +51,15 @@ def get_penetration_masks(obj_verts, obj_faces, hand_verts):
 def get_optimized_hand_fr_joints_v4_anchors(joints, base_pts, tot_base_pts_trans, tot_base_normals_trans, with_contact_opt=False, nn_hand_params=24, rt_vars=False, with_proj=False, obj_verts_trans=None, obj_faces=None, with_params_smoothing=False, dist_thres=0.005, with_ctx_mask=False):
   # obj_verts_trans, obj_faces
   joints = torch.from_numpy(joints).float() #  # joints 
-  base_pts = torch.from_numpy(base_pts).float() #  # base_pts 
+  # base_pts = torch.from_numpy(base_pts).float() #  # base_pts 
   
   if nn_hand_params < 45:
     use_pca = True
   else:
     use_pca = False
   
-  tot_base_pts_trans = torch.from_numpy(tot_base_pts_trans).float()
-  tot_base_normals_trans = torch.from_numpy(tot_base_normals_trans).float()
+  # tot_base_pts_trans = torch.from_numpy(tot_base_pts_trans).float()
+  # tot_base_normals_trans = torch.from_numpy(tot_base_normals_trans).float()
   ### start optimization ###
   # setup MANO layer
   # mano_path = "/data1/xueyi/mano_models/mano/models"
@@ -139,48 +139,48 @@ def get_optimized_hand_fr_joints_v4_anchors(joints, base_pts, tot_base_pts_trans
   # )
   
   # 
-  dist_joints_to_base_pts = torch.sum(
-    (joints.unsqueeze(-2) - tot_base_pts_trans.unsqueeze(1)) ** 2, dim=-1 # nf x nnjoints x nnbasepts #
-  )
+  # dist_joints_to_base_pts = torch.sum(
+  #   (joints.unsqueeze(-2) - tot_base_pts_trans.unsqueeze(1)) ** 2, dim=-1 # nf x nnjoints x nnbasepts #
+  # )
   
-  nn_base_pts = dist_joints_to_base_pts.size(-1)
-  nn_joints = dist_joints_to_base_pts.size(1)
+  # nn_base_pts = dist_joints_to_base_pts.size(-1)
+  # nn_joints = dist_joints_to_base_pts.size(1)
   
-  dist_joints_to_base_pts = torch.sqrt(dist_joints_to_base_pts) # nf x nnjoints x nnbasepts #
-  minn_dist, minn_dist_idx = torch.min(dist_joints_to_base_pts, dim=-1) # nf x nnjoints #
+  # dist_joints_to_base_pts = torch.sqrt(dist_joints_to_base_pts) # nf x nnjoints x nnbasepts #
+  # minn_dist, minn_dist_idx = torch.min(dist_joints_to_base_pts, dim=-1) # nf x nnjoints #
   
-  nk_contact_pts = 2
-  minn_dist[:, :-5] = 1e9
-  minn_topk_dist, minn_topk_idx = torch.topk(minn_dist, k=nk_contact_pts, largest=False) # 
-  # joints_idx_rng_exp = torch.arange(nn_joints).unsqueeze(0) == 
-  minn_topk_mask = torch.zeros_like(minn_dist)
-  # minn_topk_mask[minn_topk_idx] = 1. # nf x nnjoints #
-  minn_topk_mask[:, -5: -3] = 1.
-  basepts_idx_range = torch.arange(nn_base_pts).unsqueeze(0).unsqueeze(0)
-  minn_dist_mask = basepts_idx_range == minn_dist_idx.unsqueeze(-1) # nf x nnjoints x nnbasepts
-  # for seq 101
-  # minn_dist_mask[31:, -5, :] = minn_dist_mask[30: 31, -5, :]
-  minn_dist_mask = minn_dist_mask.float()
+  # nk_contact_pts = 2
+  # minn_dist[:, :-5] = 1e9
+  # minn_topk_dist, minn_topk_idx = torch.topk(minn_dist, k=nk_contact_pts, largest=False) # 
+  # # joints_idx_rng_exp = torch.arange(nn_joints).unsqueeze(0) == 
+  # minn_topk_mask = torch.zeros_like(minn_dist)
+  # # minn_topk_mask[minn_topk_idx] = 1. # nf x nnjoints #
+  # minn_topk_mask[:, -5: -3] = 1.
+  # basepts_idx_range = torch.arange(nn_base_pts).unsqueeze(0).unsqueeze(0)
+  # minn_dist_mask = basepts_idx_range == minn_dist_idx.unsqueeze(-1) # nf x nnjoints x nnbasepts
+  # # for seq 101
+  # # minn_dist_mask[31:, -5, :] = minn_dist_mask[30: 31, -5, :]
+  # minn_dist_mask = minn_dist_mask.float()
   
-  ## tot base pts 
-  tot_base_pts_trans_disp = torch.sum(
-    (tot_base_pts_trans[1:, :, :] - tot_base_pts_trans[:-1, :, :]) ** 2, dim=-1 # (nf - 1) x nn_base_pts displacement 
-  )
-  ### tot base pts trans disp ###
-  tot_base_pts_trans_disp = torch.sqrt(tot_base_pts_trans_disp).mean(dim=-1) # (nf - 1)
-  # tot_base_pts_trans_disp_mov_thres = 1e-20
-  tot_base_pts_trans_disp_mov_thres = 3e-4
-  tot_base_pts_trans_disp_mask = tot_base_pts_trans_disp >= tot_base_pts_trans_disp_mov_thres
-  tot_base_pts_trans_disp_mask = torch.cat(
-    [tot_base_pts_trans_disp_mask, tot_base_pts_trans_disp_mask[-1:]], dim=0
-  )
+  # ## tot base pts 
+  # tot_base_pts_trans_disp = torch.sum(
+  #   (tot_base_pts_trans[1:, :, :] - tot_base_pts_trans[:-1, :, :]) ** 2, dim=-1 # (nf - 1) x nn_base_pts displacement 
+  # )
+  # ### tot base pts trans disp ###
+  # tot_base_pts_trans_disp = torch.sqrt(tot_base_pts_trans_disp).mean(dim=-1) # (nf - 1)
+  # # tot_base_pts_trans_disp_mov_thres = 1e-20
+  # tot_base_pts_trans_disp_mov_thres = 3e-4
+  # tot_base_pts_trans_disp_mask = tot_base_pts_trans_disp >= tot_base_pts_trans_disp_mov_thres
+  # tot_base_pts_trans_disp_mask = torch.cat(
+  #   [tot_base_pts_trans_disp_mask, tot_base_pts_trans_disp_mask[-1:]], dim=0
+  # )
   
-  attraction_mask_new = (tot_base_pts_trans_disp_mask.float().unsqueeze(-1).unsqueeze(-1) + minn_dist_mask.float()) > 1.5
+  # attraction_mask_new = (tot_base_pts_trans_disp_mask.float().unsqueeze(-1).unsqueeze(-1) + minn_dist_mask.float()) > 1.5
   
   
   
-  minn_topk_mask = (minn_dist_mask + minn_topk_mask.float().unsqueeze(-1)) > 1.5
-  print(f"minn_dist_mask: {minn_dist_mask.size()}")
+  # minn_topk_mask = (minn_dist_mask + minn_topk_mask.float().unsqueeze(-1)) > 1.5
+  # print(f"minn_dist_mask: {minn_dist_mask.size()}")
   s = 1.0
   # affinity_scores = get_affinity_fr_dist(dist_joints_to_base_pts, s=s)
 
@@ -231,21 +231,21 @@ def get_optimized_hand_fr_joints_v4_anchors(joints, base_pts, tot_base_pts_trans
       print('\tRotation Smoothness Loss: {}'.format(joints_pred_loss.item()))
   
   # 
-  print(tot_base_pts_trans.size())
-  diff_base_pts_trans = torch.sum((tot_base_pts_trans[1:, :, :] - tot_base_pts_trans[:-1, :, :]) ** 2, dim=-1) # (nf - 1) x nn_base_pts
-  print(f"diff_base_pts_trans: {diff_base_pts_trans.size()}")
-  diff_base_pts_trans = diff_base_pts_trans.mean(dim=-1)
-  diff_base_pts_trans_threshold = 1e-20
-  diff_base_pts_trans_mask = diff_base_pts_trans > diff_base_pts_trans_threshold # (nf - 1) ### the mask of the tranformed base pts
-  diff_base_pts_trans_mask = diff_base_pts_trans_mask.float()
-  print(f"diff_base_pts_trans_mask: {diff_base_pts_trans_mask.size()}, diff_base_pts_trans: {diff_base_pts_trans.size()}")
-  diff_last_frame_mask = torch.tensor([0,], dtype=torch.float32).to(diff_base_pts_trans_mask.device) + diff_base_pts_trans_mask[-1]
-  diff_base_pts_trans_mask = torch.cat(
-    [diff_base_pts_trans_mask, diff_last_frame_mask], dim=0 # nf tensor
-  )
+  # print(tot_base_pts_trans.size())
+  # diff_base_pts_trans = torch.sum((tot_base_pts_trans[1:, :, :] - tot_base_pts_trans[:-1, :, :]) ** 2, dim=-1) # (nf - 1) x nn_base_pts
+  # print(f"diff_base_pts_trans: {diff_base_pts_trans.size()}")
+  # diff_base_pts_trans = diff_base_pts_trans.mean(dim=-1)
+  # diff_base_pts_trans_threshold = 1e-20
+  # diff_base_pts_trans_mask = diff_base_pts_trans > diff_base_pts_trans_threshold # (nf - 1) ### the mask of the tranformed base pts
+  # diff_base_pts_trans_mask = diff_base_pts_trans_mask.float()
+  # print(f"diff_base_pts_trans_mask: {diff_base_pts_trans_mask.size()}, diff_base_pts_trans: {diff_base_pts_trans.size()}")
+  # diff_last_frame_mask = torch.tensor([0,], dtype=torch.float32).to(diff_base_pts_trans_mask.device) + diff_base_pts_trans_mask[-1]
+  # diff_base_pts_trans_mask = torch.cat(
+  #   [diff_base_pts_trans_mask, diff_last_frame_mask], dim=0 # nf tensor
+  # )
   # attraction_mask = (diff_base_pts_trans_mask.unsqueeze(-1).unsqueeze(-1) + minn_topk_mask.float()) > 1.5
-  attraction_mask = minn_topk_mask.float()
-  attraction_mask = attraction_mask.float()
+  # attraction_mask = minn_topk_mask.float()
+  # attraction_mask = attraction_mask.float()
   
   # the direction of the normal vector and the moving direction of the object point -> whether the point should be selected
   # the contact maps of the object should be like? #
@@ -1184,7 +1184,12 @@ def reconstruct_from_file(single_seq_path):
                   tot_data[cur_k].append(cur_data[cur_k][ :clip_ending_idxes[i_tag]])
     
     for cur_k in tot_data:
-      if cur_k in ["tot_base_pts", "tot_base_normals", "tot_obj_rot", "tot_obj_transl", "tot_obj_pcs", "tot_rhand_joints", "tot_gt_rhand_joints"]:
+      print(f"cur_k: {cur_k}")
+      for aa in tot_data[cur_k]:
+        print(aa.shape)
+      if cur_k in ['tot_base_pts', 'tot_base_normals']:
+        continue
+      elif cur_k in ["tot_base_pts", "tot_base_normals", "tot_obj_rot", "tot_obj_transl", "tot_obj_pcs", "tot_rhand_joints", "tot_gt_rhand_joints"]:
           tot_data[cur_k] = np.concatenate(tot_data[cur_k], axis=1)
       else:
           tot_data[cur_k] = np.concatenate(tot_data[cur_k], axis=0)
@@ -1202,8 +1207,8 @@ def reconstruct_from_file(single_seq_path):
     
     targets = data['targets'] # # targets # #
     outputs = data['outputs'] #  
-    tot_base_pts = data["tot_base_pts"][0] # total base pts, total base normals #
-    tot_base_normals = data['tot_base_normals'][0] # nn_base_normals #
+    # tot_base_pts = data["tot_base_pts"][0] # total base pts, total base normals #
+    # tot_base_normals = data['tot_base_normals'][0] # nn_base_normals #
     
     
     
@@ -1213,22 +1218,22 @@ def reconstruct_from_file(single_seq_path):
     tot_obj_transl = data['tot_obj_transl'][0]
     print(f"tot_obj_rot: {tot_obj_rot.shape}, tot_obj_transl: {tot_obj_transl.shape}")
     
-    if len(tot_base_pts.shape) == 2:
-      # numpy array # # tot base pts #
-      tot_base_pts_trans = np.matmul(tot_base_pts.reshape(1, tot_base_pts.shape[0], 3), tot_obj_rot) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1])
-      tot_base_pts = np.matmul(tot_base_pts, tot_obj_rot[0]) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1])[0] 
+    # if len(tot_base_pts.shape) == 2:
+    #   # numpy array # # tot base pts #
+    #   tot_base_pts_trans = np.matmul(tot_base_pts.reshape(1, tot_base_pts.shape[0], 3), tot_obj_rot) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1])
+    #   tot_base_pts = np.matmul(tot_base_pts, tot_obj_rot[0]) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1])[0] 
       
-      tot_base_normals_trans = np.matmul( # # 
-        tot_base_normals.reshape(1, tot_base_normals.shape[0], 3), tot_obj_rot
-      ) 
-    else:
-      print(f"tot_base_pts: {tot_base_pts.shape}, tot_obj_rot: {tot_obj_rot.shape}, tot_obj_transl: {tot_obj_transl.shape}")
-      tot_base_pts_trans = np.matmul(tot_base_pts, tot_obj_rot) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1])
-      tot_base_pts = np.matmul(tot_base_pts, tot_obj_rot) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1])
+    #   tot_base_normals_trans = np.matmul( # # 
+    #     tot_base_normals.reshape(1, tot_base_normals.shape[0], 3), tot_obj_rot
+    #   ) 
+    # else:
+    #   print(f"tot_base_pts: {tot_base_pts.shape}, tot_obj_rot: {tot_obj_rot.shape}, tot_obj_transl: {tot_obj_transl.shape}")
+    #   tot_base_pts_trans = np.matmul(tot_base_pts, tot_obj_rot) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1])
+    #   tot_base_pts = np.matmul(tot_base_pts, tot_obj_rot) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1])
       
-      tot_base_normals_trans = np.matmul(
-        tot_base_normals, tot_obj_rot
-      )
+    #   tot_base_normals_trans = np.matmul(
+    #     tot_base_normals, tot_obj_rot
+    #   )
 
 
     
@@ -1237,7 +1242,7 @@ def reconstruct_from_file(single_seq_path):
     
     targets = np.matmul(targets, tot_obj_rot) + tot_obj_transl.reshape(tot_obj_transl.shape[0], 1, tot_obj_transl.shape[1]) # ws x nn_verts x 3 #
     # denoise relative positions 
-    print(f"tot_base_pts: {tot_base_pts.shape}")
+    # print(f"tot_base_pts: {tot_base_pts.shape}")
     
     
     #### obj_verts_trans, obj_faces ####
@@ -1264,7 +1269,7 @@ def reconstruct_from_file(single_seq_path):
     with_contact_opt = True
     with_ctx_mask = False
 
-    bf_ct_optimized_dict, bf_proj_optimized_dict, optimized_dict = get_optimized_hand_fr_joints_v4_anchors(outputs, tot_base_pts, tot_base_pts_trans, tot_base_normals_trans, with_contact_opt=with_contact_opt, nn_hand_params=nn_hand_params, rt_vars=True, with_proj=with_proj, obj_verts_trans=obj_verts_trans, obj_faces=obj_faces, with_params_smoothing=with_params_smoothing, dist_thres=dist_thres, with_ctx_mask=with_ctx_mask)
+    bf_ct_optimized_dict, bf_proj_optimized_dict, optimized_dict = get_optimized_hand_fr_joints_v4_anchors(outputs, None, None, None, with_contact_opt=with_contact_opt, nn_hand_params=nn_hand_params, rt_vars=True, with_proj=with_proj, obj_verts_trans=obj_verts_trans, obj_faces=obj_faces, with_params_smoothing=with_params_smoothing, dist_thres=dist_thres, with_ctx_mask=with_ctx_mask)
 
     
 
@@ -1274,12 +1279,12 @@ def reconstruct_from_file(single_seq_path):
     optimized_sv_infos.update(bf_ct_optimized_dict)
     optimized_sv_infos.update(bf_proj_optimized_dict)
     optimized_sv_infos.update(optimized_dict)
-    optimized_sv_infos.update(
-      {
-        'tot_base_pts_trans': tot_base_pts_trans,
-        'tot_base_normals_trans': tot_base_normals_trans
-      }
-    )
+    # optimized_sv_infos.update(
+    #   {
+    #     'tot_base_pts_trans': tot_base_pts_trans,
+    #     'tot_base_normals_trans': tot_base_normals_trans
+    #   }
+    # )
     
     
     optimized_sv_infos.update({'predicted_info': data})