You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
157 lines
5.7 KiB
157 lines
5.7 KiB
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from mmpose.apis import (get_track_id, inference_bottom_up_pose_model,
|
|
inference_top_down_pose_model, init_pose_model,
|
|
vis_pose_tracking_result)
|
|
from mmpose.datasets.dataset_info import DatasetInfo
|
|
|
|
|
|
def test_top_down_pose_tracking_demo():
|
|
# COCO demo
|
|
# build the pose model from a config file and a checkpoint file
|
|
pose_model = init_pose_model(
|
|
'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/'
|
|
'coco/res50_coco_256x192.py',
|
|
None,
|
|
device='cpu')
|
|
image_name = 'tests/data/coco/000000000785.jpg'
|
|
dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])
|
|
person_result = [{'bbox': [50, 50, 50, 100]}]
|
|
|
|
# test a single image, with a list of bboxes.
|
|
pose_results, _ = inference_top_down_pose_model(
|
|
pose_model,
|
|
image_name,
|
|
person_result,
|
|
format='xywh',
|
|
dataset_info=dataset_info)
|
|
pose_results, next_id = get_track_id(pose_results, [], next_id=0)
|
|
# show the results
|
|
vis_pose_tracking_result(
|
|
pose_model, image_name, pose_results, dataset_info=dataset_info)
|
|
pose_results_last = pose_results
|
|
|
|
# AIC demo
|
|
pose_model = init_pose_model(
|
|
'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/'
|
|
'aic/res50_aic_256x192.py',
|
|
None,
|
|
device='cpu')
|
|
image_name = 'tests/data/aic/054d9ce9201beffc76e5ff2169d2af2f027002ca.jpg'
|
|
dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])
|
|
# test a single image, with a list of bboxes.
|
|
pose_results, _ = inference_top_down_pose_model(
|
|
pose_model,
|
|
image_name,
|
|
person_result,
|
|
format='xywh',
|
|
dataset_info=dataset_info)
|
|
pose_results, next_id = get_track_id(pose_results, pose_results_last,
|
|
next_id)
|
|
for pose_result in pose_results:
|
|
del pose_result['bbox']
|
|
pose_results, next_id = get_track_id(pose_results, pose_results_last,
|
|
next_id)
|
|
|
|
# show the results
|
|
vis_pose_tracking_result(
|
|
pose_model, image_name, pose_results, dataset_info=dataset_info)
|
|
|
|
# OneHand10K demo
|
|
# build the pose model from a config file and a checkpoint file
|
|
pose_model = init_pose_model(
|
|
'configs/hand/2d_kpt_sview_rgb_img/topdown_heatmap/'
|
|
'onehand10k/res50_onehand10k_256x256.py',
|
|
None,
|
|
device='cpu')
|
|
image_name = 'tests/data/onehand10k/9.jpg'
|
|
dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])
|
|
# test a single image, with a list of bboxes.
|
|
pose_results, _ = inference_top_down_pose_model(
|
|
pose_model,
|
|
image_name, [{
|
|
'bbox': [10, 10, 30, 30]
|
|
}],
|
|
format='xywh',
|
|
dataset_info=dataset_info)
|
|
pose_results, next_id = get_track_id(pose_results, pose_results_last,
|
|
next_id)
|
|
# show the results
|
|
vis_pose_tracking_result(
|
|
pose_model, image_name, pose_results, dataset_info=dataset_info)
|
|
|
|
# InterHand2D demo
|
|
pose_model = init_pose_model(
|
|
'configs/hand/2d_kpt_sview_rgb_img/topdown_heatmap/'
|
|
'interhand2d/res50_interhand2d_all_256x256.py',
|
|
None,
|
|
device='cpu')
|
|
image_name = 'tests/data/interhand2.6m/image2017.jpg'
|
|
dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])
|
|
# test a single image, with a list of bboxes.
|
|
pose_results, _ = inference_top_down_pose_model(
|
|
pose_model,
|
|
image_name, [{
|
|
'bbox': [50, 50, 0, 0]
|
|
}],
|
|
format='xywh',
|
|
dataset_info=dataset_info)
|
|
pose_results, next_id = get_track_id(pose_results, [], next_id=0)
|
|
# show the results
|
|
vis_pose_tracking_result(
|
|
pose_model, image_name, pose_results, dataset_info=dataset_info)
|
|
pose_results_last = pose_results
|
|
|
|
# MPII demo
|
|
pose_model = init_pose_model(
|
|
'configs/body/2d_kpt_sview_rgb_img/topdown_heatmap/'
|
|
'mpii/res50_mpii_256x256.py',
|
|
None,
|
|
device='cpu')
|
|
image_name = 'tests/data/mpii/004645041.jpg'
|
|
dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])
|
|
# test a single image, with a list of bboxes.
|
|
pose_results, _ = inference_top_down_pose_model(
|
|
pose_model,
|
|
image_name, [{
|
|
'bbox': [50, 50, 0, 0]
|
|
}],
|
|
format='xywh',
|
|
dataset_info=dataset_info)
|
|
pose_results, next_id = get_track_id(pose_results, pose_results_last,
|
|
next_id)
|
|
# show the results
|
|
vis_pose_tracking_result(
|
|
pose_model, image_name, pose_results, dataset_info=dataset_info)
|
|
|
|
|
|
def test_bottom_up_pose_tracking_demo():
|
|
# COCO demo
|
|
# build the pose model from a config file and a checkpoint file
|
|
pose_model = init_pose_model(
|
|
'configs/body/2d_kpt_sview_rgb_img/associative_embedding/'
|
|
'coco/res50_coco_512x512.py',
|
|
None,
|
|
device='cpu')
|
|
|
|
image_name = 'tests/data/coco/000000000785.jpg'
|
|
dataset_info = DatasetInfo(pose_model.cfg.data['test']['dataset_info'])
|
|
|
|
pose_results, _ = inference_bottom_up_pose_model(
|
|
pose_model, image_name, dataset_info=dataset_info)
|
|
|
|
pose_results, next_id = get_track_id(pose_results, [], next_id=0)
|
|
|
|
# show the results
|
|
vis_pose_tracking_result(
|
|
pose_model, image_name, pose_results, dataset_info=dataset_info)
|
|
|
|
pose_results_last = pose_results
|
|
|
|
# oks
|
|
pose_results, next_id = get_track_id(
|
|
pose_results, pose_results_last, next_id=next_id, use_oks=True)
|
|
|
|
pose_results_last = pose_results
|
|
# one_euro
|
|
pose_results, next_id = get_track_id(
|
|
pose_results, pose_results_last, next_id=next_id, use_one_euro=True)
|
|
|