Merge pull request #65 from kaseris/fix/visualization-args

Remove the hard coding of the model checkpoint

tools/visualize_skel_movement.py

@@ -13,6 +13,8 @@
 argparser.add_argument('--dataset', type=str, required=True, help='Path to the dataset.')
 argparser.add_argument('--sample', type=int, required=True, help='Sample index to visualize.')
 argparser.add_argument('--cache-file', type=str, required=False, help='Path to the cache file.')
+argparser.add_argument('--checkpoint', type=str, required=False, help='Path to the checkpoint file.',
+ default='/home/kaseris/Documents/mount/checkpoints_forecasting/presto-class/checkpoint_epoch_16_2024-01-05_092620.pt')
 
 args = argparser.parse_args()
 
@@ -30,7 +32,7 @@
  # checkpoint = torch.load('/home/kaseris/Documents/mount/checkpoints_forecasting/presto-class/checkpoint_epoch_16_2024-01-05_092620.pt')
  # model_state_dict = checkpoint['model_state_dict']
  # model.load_state_dict(model_state_dict)
- model.from_pretrained('/home/kaseris/Documents/mount/checkpoints_forecasting/presto-class/checkpoint_epoch_16_2024-01-05_092620.pt')
+ model.from_pretrained(args.checkpoint)
  model = model.to('cpu')
  skeleton, label = dataset[args.sample]
  seq_len, n_bodies, n_joints, n_dims = skeleton.shape