167 lines
6.3 KiB
Python
167 lines
6.3 KiB
Python
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()
|