| 81 | return det_results, max_person, valid_frame_idx,img_list,frame_id |
| 82 | |
| 83 | def visualize(results,draw_num,save_path,detection_all=None,input_type='image'): |
| 84 | pred_mesh = results['pred_mesh']#*1000 # (N*T, V, 3) |
| 85 | focal_l = results['focal_l'] |
| 86 | center_pt = results['center_pt'] |
| 87 | if input_type=='video': |
| 88 | img_idx_list = results['img_idx'] |
| 89 | img_path_list = results['img_path_video'] |
| 90 | max_person = results['max_person'] |
| 91 | video_name = results['video_name'] |
| 92 | videowriter = imageio.get_writer(osp.join(save_path, f"{video_name}_results_in_2d.mp4"), fps=results['fps']) |
| 93 | frame_num = max(img_idx_list) |
| 94 | for i in tqdm(range(frame_num)): |
| 95 | chosen_mask = img_idx_list == i |
| 96 | chosen_idx = np.where(chosen_mask>0)[0] |
| 97 | if chosen_idx.shape[0] == 0: |
| 98 | continue |
| 99 | pred_mesh_T = pred_mesh[chosen_idx] |
| 100 | focal_T = focal_l[chosen_idx[0]] |
| 101 | center_pt_T = center_pt[chosen_idx[0]] |
| 102 | img_path = img_path_list[chosen_idx[0]] |
| 103 | img = cv2.imread(img_path) |
| 104 | ori_img_height, ori_img_width = img.shape[:2] |
| 105 | rgb, depth = render_mesh(ori_img_height, ori_img_width, pred_mesh_T/1000.0, smpl.faces, {'focal': focal_T, 'princpt': center_pt_T}) |
| 106 | valid_mask = (depth > 0)[:,:,None] |
| 107 | rendered_img = rgb * valid_mask + img[:,:,::-1] * (1-valid_mask) |
| 108 | cv2.imwrite(osp.join(save_path, f"{video_name}_results_in_2d.jpg"), rendered_img.astype(np.uint8)[...,::-1]) |
| 109 | videowriter.append_data(rendered_img.astype(np.uint8)) |
| 110 | videowriter.close() |
| 111 | else: |
| 112 | img_path_list = results['img_path'] |
| 113 | img_name = results['img_name'] |
| 114 | idx_list =results['idx'] |
| 115 | for mesh_idx in tqdm(range(pred_mesh.shape[0])): |
| 116 | img_path = img_path_list[mesh_idx] |
| 117 | img = cv2.imread(img_path) |
| 118 | ori_img_height, ori_img_width = img.shape[:2] |
| 119 | dirpath = osp.join(save_path,img_name[mesh_idx],str(idx_list[mesh_idx])) |
| 120 | if os.path.exists(dirpath) and os.path.isdir(dirpath): |
| 121 | shutil.rmtree(dirpath) |
| 122 | for i in range(draw_num): |
| 123 | pred_mesh_T = pred_mesh[mesh_idx,i] |
| 124 | rgb, depth = render_mesh(ori_img_height, ori_img_width, [pred_mesh_T/1000], smpl.faces, {'focal': focal_l[mesh_idx], 'princpt': center_pt[mesh_idx]}) |
| 125 | |
| 126 | valid_mask = (depth > 0)[:,:,None] |
| 127 | rendered_img = rgb * valid_mask + img[:,:,::-1] * (1-valid_mask) |
| 128 | save_path_img = osp.join(save_path,img_name[mesh_idx],str(idx_list[mesh_idx]),f"idx{idx_list[mesh_idx]} hypo{i}.jpg") |
| 129 | os.makedirs(osp.join(save_path,img_name[mesh_idx],str(idx_list[mesh_idx])),exist_ok=True) |
| 130 | cv2.imwrite(save_path_img, rendered_img.astype(np.uint8)[...,::-1]) |
| 131 | |
| 132 | class DeltaDepth(nn.Module): |
| 133 | def __init__(self,bs): |