(file_name, result_path)
| 197 | return img |
| 198 | |
| 199 | def process_single_video(file_name, result_path): |
| 200 | video_name = os.path.basename(file_name).split('.pkl')[0] |
| 201 | |
| 202 | with open(file_name, 'rb') as f: |
| 203 | results_dict = pickle.load(f) |
| 204 | |
| 205 | focals = results_dict['focal'] |
| 206 | princpts = results_dict['princpt'] |
| 207 | all_pose = results_dict['smplx'] |
| 208 | height, width = results_dict['height'], results_dict['width'] |
| 209 | all_pose = torch.tensor(all_pose).cuda() |
| 210 | |
| 211 | g, b, l, r, j, s, exp, cam_trans = \ |
| 212 | all_pose[:, :3], all_pose[:, 3:66], all_pose[:, 66:111], all_pose[:, 111:156], all_pose[:, 156:159], \ |
| 213 | all_pose[:, 159:169], all_pose[:, 169:179], all_pose[:, 179:182] |
| 214 | assert len(all_pose[0]) == 182 |
| 215 | |
| 216 | meshes = get_coord(g, b, l, r, j, s, exp, cam_trans[0][None], mesh=True).cpu().numpy() |
| 217 | |
| 218 | if os.path.isdir(args.pkl_file_path): |
| 219 | bar = enumerate(results_dict['total_valid_index']) |
| 220 | else: |
| 221 | bar = enumerate(tqdm(results_dict['total_valid_index'])) |
| 222 | |
| 223 | if args.overlay: |
| 224 | if args.video_path is None: |
| 225 | raw_img_list = get_img_list(video_name, from_url=True) |
| 226 | else: |
| 227 | raw_img_list = get_img_list(video_name) |
| 228 | |
| 229 | img_list = [] |
| 230 | text = text_dict[video_name.replace('/', '')] |
| 231 | for idx, index in bar: |
| 232 | |
| 233 | if args.overlay: |
| 234 | # render overlay |
| 235 | raw_img = raw_img_list[index] |
| 236 | img = render(raw_img.copy(), meshes[idx], smpl_x.face, {'focal': focals[0], 'princpt': princpts[0]}) |
| 237 | img = np.array(np.concatenate((raw_img,img), axis=1), dtype=np.uint8) |
| 238 | size = (2*width,height) |
| 239 | else: |
| 240 | # render with background |
| 241 | img = render(background, meshes[idx], smpl_x.face, {'focal': pred_focals, 'princpt': pred_princpts}).astype(np.uint8) |
| 242 | size = (predefined_width,predefined_height) |
| 243 | img_list.append(img) |
| 244 | |
| 245 | put_text_with_newline(img, text, org, font, font_scale, color, thickness) |
| 246 | |
| 247 | out = cv2.VideoWriter(result_path + f'/{video_name}.mp4', 0x7634706d, 24, size) |
| 248 | for idx in range(len(img_list)): |
| 249 | out.write(img_list[idx]) |
| 250 | out.release() |
| 251 | |
| 252 | if __name__ == "__main__": |
| 253 | parser = ArgumentParser() |
no test coverage detected