Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Vendored
+166
@@ -0,0 +1,166 @@
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
from pathlib import Path
|
||||
import traceback
|
||||
from typing import List, Optional
|
||||
|
||||
import pandas as pd
|
||||
import torch
|
||||
from filelock import FileLock
|
||||
from hamer.configs import dataset_eval_config
|
||||
from hamer.datasets import create_dataset
|
||||
from hamer.utils import Evaluator, recursive_to
|
||||
from tqdm import tqdm
|
||||
|
||||
from hamer.configs import CACHE_DIR_HAMER
|
||||
from hamer.models import HAMER, download_models, load_hamer, DEFAULT_CHECKPOINT
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Evaluate trained models')
|
||||
parser.add_argument('--checkpoint', type=str, default=DEFAULT_CHECKPOINT, help='Path to pretrained model checkpoint')
|
||||
parser.add_argument('--results_folder', type=str, default='results', help='Path to results folder.')
|
||||
parser.add_argument('--dataset', type=str, default='FREIHAND-VAL,HO3D-VAL,NEWDAYS-TEST-ALL,NEWDAYS-TEST-VIS,NEWDAYS-TEST-OCC,EPICK-TEST-ALL,EPICK-TEST-VIS,EPICK-TEST-OCC,EGO4D-TEST-ALL,EGO4D-TEST-VIS,EGO4D-TEST-OCC', help='Dataset to evaluate')
|
||||
parser.add_argument('--batch_size', type=int, default=16, help='Batch size for inference')
|
||||
parser.add_argument('--num_samples', type=int, default=1, help='Number of test samples to draw')
|
||||
parser.add_argument('--num_workers', type=int, default=8, help='Number of workers used for data loading')
|
||||
parser.add_argument('--log_freq', type=int, default=10, help='How often to log results')
|
||||
parser.add_argument('--shuffle', dest='shuffle', action='store_true', default=False, help='Shuffle the dataset during evaluation')
|
||||
parser.add_argument('--exp_name', type=str, default=None, help='Experiment name')
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Download and load checkpoints
|
||||
download_models(CACHE_DIR_HAMER)
|
||||
model, model_cfg = load_hamer(args.checkpoint)
|
||||
|
||||
# Setup HMR2.0 model
|
||||
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
|
||||
# Load config and run eval, one dataset at a time
|
||||
print('Evaluating on datasets: {}'.format(args.dataset), flush=True)
|
||||
for dataset in args.dataset.split(','):
|
||||
dataset_cfg = dataset_eval_config()[dataset]
|
||||
args.dataset = dataset
|
||||
run_eval(model, model_cfg, dataset_cfg, device, args)
|
||||
|
||||
def run_eval(model, model_cfg, dataset_cfg, device, args):
|
||||
|
||||
# List of metrics to log
|
||||
if args.dataset in ['FREIHAND-VAL', 'HO3D-VAL']:
|
||||
metrics = None
|
||||
preds = ['vertices', 'keypoints_3d']
|
||||
pck_thresholds = None
|
||||
rescale_factor = -1
|
||||
elif args.dataset in ['NEWDAYS-TEST-ALL', 'NEWDAYS-TEST-VIS', 'NEWDAYS-TEST-OCC',
|
||||
'EPICK-TEST-ALL', 'EPICK-TEST-VIS', 'EPICK-TEST-OCC',
|
||||
'EGO4D-TEST-ALL', 'EGO4D-TEST-VIS', 'EGO4D-TEST-OCC']:
|
||||
metrics = ['mode_kpl2']
|
||||
preds = None
|
||||
pck_thresholds = [0.05, 0.1, 0.15]
|
||||
rescale_factor = 2
|
||||
|
||||
# Create dataset and data loader
|
||||
dataset = create_dataset(model_cfg, dataset_cfg, train=False, rescale_factor=rescale_factor)
|
||||
dataloader = torch.utils.data.DataLoader(dataset, args.batch_size, shuffle=args.shuffle, num_workers=args.num_workers)
|
||||
|
||||
# Setup evaluator object
|
||||
evaluator = Evaluator(
|
||||
dataset_length=dataset.__len__(),
|
||||
dataset=args.dataset,
|
||||
keypoint_list=dataset_cfg.KEYPOINT_LIST,
|
||||
pelvis_ind=model_cfg.EXTRA.PELVIS_IND,
|
||||
metrics=metrics,
|
||||
preds=preds,
|
||||
pck_thresholds=pck_thresholds,
|
||||
)
|
||||
|
||||
# Go over the images in the dataset.
|
||||
try:
|
||||
for i, batch in enumerate(tqdm(dataloader)):
|
||||
batch = recursive_to(batch, device)
|
||||
with torch.no_grad():
|
||||
out = model(batch)
|
||||
evaluator(out, batch)
|
||||
if i % args.log_freq == args.log_freq - 1:
|
||||
evaluator.log()
|
||||
evaluator.log()
|
||||
error = None
|
||||
except (Exception, KeyboardInterrupt) as e:
|
||||
traceback.print_exc()
|
||||
error = repr(e)
|
||||
i = 0
|
||||
|
||||
# Append results to file
|
||||
if metrics is not None:
|
||||
metrics_dict = evaluator.get_metrics_dict()
|
||||
results_csv = os.path.join(args.results_folder, 'eval_regression.csv')
|
||||
save_eval_result(results_csv, metrics_dict, args.checkpoint, args.dataset, error=error, iters_done=i, exp_name=args.exp_name)
|
||||
if preds is not None:
|
||||
results_json = os.path.join(args.results_folder, '%s.json' % args.dataset.lower())
|
||||
preds_dict = evaluator.get_preds_dict()
|
||||
save_preds_result(results_json, preds_dict)
|
||||
|
||||
def save_eval_result(
|
||||
csv_path: str,
|
||||
metric_dict: float,
|
||||
checkpoint_path: str,
|
||||
dataset_name: str,
|
||||
# start_time: pd.Timestamp,
|
||||
error: Optional[str] = None,
|
||||
iters_done=None,
|
||||
exp_name=None,
|
||||
) -> None:
|
||||
"""Save evaluation results for a single scene file to a common CSV file."""
|
||||
|
||||
timestamp = pd.Timestamp.now()
|
||||
exists: bool = os.path.exists(csv_path)
|
||||
exp_name = exp_name or Path(checkpoint_path).parent.parent.name
|
||||
|
||||
# save each metric as different row to the csv path
|
||||
metric_names = list(metric_dict.keys())
|
||||
metric_values = list(metric_dict.values())
|
||||
N = len(metric_names)
|
||||
df = pd.DataFrame(
|
||||
dict(
|
||||
timestamp=[timestamp] * N,
|
||||
checkpoint_path=[checkpoint_path] * N,
|
||||
exp_name=[exp_name] * N,
|
||||
dataset=[dataset_name] * N,
|
||||
metric_name=metric_names,
|
||||
metric_value=metric_values,
|
||||
error=[error] * N,
|
||||
iters_done=[iters_done] * N,
|
||||
),
|
||||
index=list(range(N)),
|
||||
)
|
||||
|
||||
# Lock the file to prevent multiple processes from writing to it at the same time.
|
||||
lock = FileLock(f"{csv_path}.lock", timeout=10)
|
||||
with lock:
|
||||
df.to_csv(csv_path, mode="a", header=not exists, index=False)
|
||||
|
||||
def save_preds_result(
|
||||
pred_out_path: str,
|
||||
preds_dict: float,
|
||||
) -> None:
|
||||
""" Save predictions into a json file. """
|
||||
xyz_pred_list = preds_dict['keypoints_3d']
|
||||
verts_pred_list = preds_dict['vertices']
|
||||
# make sure its only lists
|
||||
xyz_pred_list = [x.tolist() for x in xyz_pred_list]
|
||||
verts_pred_list = [x.tolist() for x in verts_pred_list]
|
||||
|
||||
# save to a json
|
||||
with open(pred_out_path, 'w') as fo:
|
||||
json.dump(
|
||||
[
|
||||
xyz_pred_list,
|
||||
verts_pred_list
|
||||
], fo)
|
||||
print('Dumped %d joints and %d verts predictions to %s' % (len(xyz_pred_list), len(verts_pred_list), pred_out_path))
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user