Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
+184
@@ -0,0 +1,184 @@
|
||||
# Copyright (c) OpenMMLab. All rights reserved.
|
||||
import argparse
|
||||
import os
|
||||
import os.path as osp
|
||||
import warnings
|
||||
|
||||
import mmcv
|
||||
import torch
|
||||
from mmcv import Config, DictAction
|
||||
from mmcv.cnn import fuse_conv_bn
|
||||
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
|
||||
from mmcv.runner import get_dist_info, init_dist, load_checkpoint
|
||||
|
||||
from mmpose.apis import multi_gpu_test, single_gpu_test
|
||||
from mmpose.datasets import build_dataloader, build_dataset
|
||||
from mmpose.models import build_posenet
|
||||
from mmpose.utils import setup_multi_processes
|
||||
|
||||
try:
|
||||
from mmcv.runner import wrap_fp16_model
|
||||
except ImportError:
|
||||
warnings.warn('auto_fp16 from mmpose will be deprecated from v0.15.0'
|
||||
'Please install mmcv>=1.1.4')
|
||||
from mmpose.core import wrap_fp16_model
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='mmpose test model')
|
||||
parser.add_argument('config', help='test config file path')
|
||||
parser.add_argument('checkpoint', help='checkpoint file')
|
||||
parser.add_argument('--out', help='output result file')
|
||||
parser.add_argument(
|
||||
'--work-dir', help='the dir to save evaluation results')
|
||||
parser.add_argument(
|
||||
'--fuse-conv-bn',
|
||||
action='store_true',
|
||||
help='Whether to fuse conv and bn, this will slightly increase'
|
||||
'the inference speed')
|
||||
parser.add_argument(
|
||||
'--gpu-id',
|
||||
type=int,
|
||||
default=0,
|
||||
help='id of gpu to use '
|
||||
'(only applicable to non-distributed testing)')
|
||||
parser.add_argument(
|
||||
'--eval',
|
||||
default=None,
|
||||
nargs='+',
|
||||
help='evaluation metric, which depends on the dataset,'
|
||||
' e.g., "mAP" for MSCOCO')
|
||||
parser.add_argument(
|
||||
'--gpu_collect',
|
||||
action='store_true',
|
||||
help='whether to use gpu to collect results')
|
||||
parser.add_argument('--tmpdir', help='tmp dir for writing some results')
|
||||
parser.add_argument(
|
||||
'--cfg-options',
|
||||
nargs='+',
|
||||
action=DictAction,
|
||||
default={},
|
||||
help='override some settings in the used config, the key-value pair '
|
||||
'in xxx=yyy format will be merged into config file. For example, '
|
||||
"'--cfg-options model.backbone.depth=18 model.backbone.with_cp=True'")
|
||||
parser.add_argument(
|
||||
'--launcher',
|
||||
choices=['none', 'pytorch', 'slurm', 'mpi'],
|
||||
default='none',
|
||||
help='job launcher')
|
||||
parser.add_argument('--local_rank', type=int, default=0)
|
||||
args = parser.parse_args()
|
||||
if 'LOCAL_RANK' not in os.environ:
|
||||
os.environ['LOCAL_RANK'] = str(args.local_rank)
|
||||
return args
|
||||
|
||||
|
||||
def merge_configs(cfg1, cfg2):
|
||||
# Merge cfg2 into cfg1
|
||||
# Overwrite cfg1 if repeated, ignore if value is None.
|
||||
cfg1 = {} if cfg1 is None else cfg1.copy()
|
||||
cfg2 = {} if cfg2 is None else cfg2
|
||||
for k, v in cfg2.items():
|
||||
if v:
|
||||
cfg1[k] = v
|
||||
return cfg1
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
cfg = Config.fromfile(args.config)
|
||||
|
||||
if args.cfg_options is not None:
|
||||
cfg.merge_from_dict(args.cfg_options)
|
||||
|
||||
# set multi-process settings
|
||||
setup_multi_processes(cfg)
|
||||
|
||||
# set cudnn_benchmark
|
||||
if cfg.get('cudnn_benchmark', False):
|
||||
torch.backends.cudnn.benchmark = True
|
||||
cfg.model.pretrained = None
|
||||
cfg.data.test.test_mode = True
|
||||
|
||||
# work_dir is determined in this priority: CLI > segment in file > filename
|
||||
if args.work_dir is not None:
|
||||
# update configs according to CLI args if args.work_dir is not None
|
||||
cfg.work_dir = args.work_dir
|
||||
elif cfg.get('work_dir', None) is None:
|
||||
# use config filename as default work_dir if cfg.work_dir is None
|
||||
cfg.work_dir = osp.join('./work_dirs',
|
||||
osp.splitext(osp.basename(args.config))[0])
|
||||
|
||||
mmcv.mkdir_or_exist(osp.abspath(cfg.work_dir))
|
||||
|
||||
# init distributed env first, since logger depends on the dist info.
|
||||
if args.launcher == 'none':
|
||||
distributed = False
|
||||
else:
|
||||
distributed = True
|
||||
init_dist(args.launcher, **cfg.dist_params)
|
||||
|
||||
# build the dataloader
|
||||
dataset = build_dataset(cfg.data.test, dict(test_mode=True))
|
||||
# step 1: give default values and override (if exist) from cfg.data
|
||||
loader_cfg = {
|
||||
**dict(seed=cfg.get('seed'), drop_last=False, dist=distributed),
|
||||
**({} if torch.__version__ != 'parrots' else dict(
|
||||
prefetch_num=2,
|
||||
pin_memory=False,
|
||||
)),
|
||||
**dict((k, cfg.data[k]) for k in [
|
||||
'seed',
|
||||
'prefetch_num',
|
||||
'pin_memory',
|
||||
'persistent_workers',
|
||||
] if k in cfg.data)
|
||||
}
|
||||
# step2: cfg.data.test_dataloader has higher priority
|
||||
test_loader_cfg = {
|
||||
**loader_cfg,
|
||||
**dict(shuffle=False, drop_last=False),
|
||||
**dict(workers_per_gpu=cfg.data.get('workers_per_gpu', 1)),
|
||||
**dict(samples_per_gpu=cfg.data.get('samples_per_gpu', 1)),
|
||||
**cfg.data.get('test_dataloader', {})
|
||||
}
|
||||
data_loader = build_dataloader(dataset, **test_loader_cfg)
|
||||
|
||||
# build the model and load checkpoint
|
||||
model = build_posenet(cfg.model)
|
||||
fp16_cfg = cfg.get('fp16', None)
|
||||
if fp16_cfg is not None:
|
||||
wrap_fp16_model(model)
|
||||
load_checkpoint(model, args.checkpoint, map_location='cpu')
|
||||
|
||||
if args.fuse_conv_bn:
|
||||
model = fuse_conv_bn(model)
|
||||
|
||||
if not distributed:
|
||||
model = MMDataParallel(model, device_ids=[args.gpu_id])
|
||||
outputs = single_gpu_test(model, data_loader)
|
||||
else:
|
||||
model = MMDistributedDataParallel(
|
||||
model.cuda(),
|
||||
device_ids=[torch.cuda.current_device()],
|
||||
broadcast_buffers=False)
|
||||
outputs = multi_gpu_test(model, data_loader, args.tmpdir,
|
||||
args.gpu_collect)
|
||||
|
||||
rank, _ = get_dist_info()
|
||||
eval_config = cfg.get('evaluation', {})
|
||||
eval_config = merge_configs(eval_config, dict(metric=args.eval))
|
||||
|
||||
if rank == 0:
|
||||
if args.out:
|
||||
print(f'\nwriting results to {args.out}')
|
||||
mmcv.dump(outputs, args.out)
|
||||
|
||||
results = dataset.evaluate(outputs, cfg.work_dir, **eval_config)
|
||||
for k, v in sorted(results.items()):
|
||||
print(f'{k}: {v}')
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user