(path, output_dir, type='image')
| 263 | plt.savefig(output_dir_3D + str(('%04d'% i)) + '_3D.png', dpi=200, format='png', bbox_inches = 'tight') |
| 264 | plt.close(fig) |
| 265 | def get_pose3D(path, output_dir, type='image'): |
| 266 | |
| 267 | # args, _ = argparse.ArgumentParser().parse_known_args() |
| 268 | # args.layers, args.channel, args.d_hid, args.frames = 3, 512, 1024, 1 # frames = 1 |
| 269 | # args.pad = (args.frames - 1) // 2 |
| 270 | # args.previous_dir = './pre_trained_model' |
| 271 | # args.n_joints, args.out_joints = 17, 17 |
| 272 | # args.type = type |
| 273 | |
| 274 | ## Reload |
| 275 | device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 276 | |
| 277 | model = {} |
| 278 | model['CFM'] = CFM(args).to(device) |
| 279 | |
| 280 | # if args.reload: |
| 281 | model_dict = model['CFM'].state_dict() |
| 282 | model_path = resolve_weights_path(args.model_weights_path, args.model_type) |
| 283 | |
| 284 | print(f"Loading weights from: {model_path}") |
| 285 | pre_dict = torch.load(model_path, map_location=device, weights_only=True) |
| 286 | for name, key in model_dict.items(): |
| 287 | model_dict[name] = pre_dict[name] |
| 288 | model['CFM'].load_state_dict(model_dict) |
| 289 | print("Load model Successfully!") |
| 290 | |
| 291 | model = model['CFM'].eval() |
| 292 | |
| 293 | ## input |
| 294 | keypoints = np.load(output_dir + 'input_2D/keypoints.npz', allow_pickle=True)['reconstruction'] |
| 295 | |
| 296 | ## 3D |
| 297 | print('\nGenerating 3D pose...') |
| 298 | |
| 299 | if type=="image": |
| 300 | i = 0 |
| 301 | img = cv2.imread(path) |
| 302 | get_3D_pose_from_image(args, keypoints, i, img, model, output_dir) |
| 303 | |
| 304 | if type=="video": |
| 305 | cap = cv2.VideoCapture(path) |
| 306 | video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) |
| 307 | for i in tqdm(range(video_length)): |
| 308 | ret, img = cap.read() |
| 309 | get_3D_pose_from_image(args, keypoints, i, img, model, output_dir) |
| 310 | |
| 311 | output_dir_2D = output_dir +'pose2D/' |
| 312 | output_dir_3D = output_dir +'pose3D/' |
| 313 | |
| 314 | print('Generating 3D pose successful!') |
| 315 | |
| 316 | ## all |
| 317 | image_dir = 'results/' |
| 318 | image_2d_dir = sorted(glob.glob(os.path.join(output_dir_2D, '*.png'))) |
| 319 | image_3d_dir = sorted(glob.glob(os.path.join(output_dir_3D, '*.png'))) |
| 320 | |
| 321 | print('\nGenerating demo...') |
| 322 | for i in tqdm(range(len(image_2d_dir))): |
no test coverage detected