liyang/hand-motion-pipeline
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
liyang
2026-09-15 14:13:49 +08:00
commit 7e4ef6f98b
7002 changed files with 1452249 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
third_party/hamer/third-party/ViTPose/mmpose/apis/__init__.py
Vendored
+20
View File
@@ -0,0 +1,20 @@
# Copyright (c) OpenMMLab. All rights reserved.
from .inference import (inference_bottom_up_pose_model,
inference_top_down_pose_model, init_pose_model,
process_mmdet_results, vis_pose_result)
from .inference_3d import (extract_pose_sequence, inference_interhand_3d_model,
inference_mesh_model, inference_pose_lifter_model,
vis_3d_mesh_result, vis_3d_pose_result)
from .inference_tracking import get_track_id, vis_pose_tracking_result
from .test import multi_gpu_test, single_gpu_test
from .train import init_random_seed, train_model
__all__ = [
'train_model', 'init_pose_model', 'inference_top_down_pose_model',
'inference_bottom_up_pose_model', 'multi_gpu_test', 'single_gpu_test',
'vis_pose_result', 'get_track_id', 'vis_pose_tracking_result',
'inference_pose_lifter_model', 'vis_3d_pose_result',
'inference_interhand_3d_model', 'extract_pose_sequence',
'inference_mesh_model', 'vis_3d_mesh_result', 'process_mmdet_results',
'init_random_seed'
]
third_party/hamer/third-party/ViTPose/mmpose/apis/inference.py
Vendored
+833
View File
@@ -0,0 +1,833 @@
# Copyright (c) OpenMMLab. All rights reserved.
import os
import warnings
import mmcv
import numpy as np
import torch
from mmcv.parallel import collate, scatter
from mmcv.runner import load_checkpoint
from PIL import Image
from mmpose.core.post_processing import oks_nms
from mmpose.datasets.dataset_info import DatasetInfo
from mmpose.datasets.pipelines import Compose
from mmpose.models import build_posenet
from mmpose.utils.hooks import OutputHook
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
def init_pose_model(config, checkpoint=None, device='cuda:0'):
"""Initialize a pose model from config file.
Args:
config (str or :obj:`mmcv.Config`): Config file path or the config
object.
checkpoint (str, optional): Checkpoint path. If left as None, the model
will not load any weights.
Returns:
nn.Module: The constructed detector.
"""
if isinstance(config, str):
config = mmcv.Config.fromfile(config)
elif not isinstance(config, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
f'but got {type(config)}')
config.model.pretrained = None
model = build_posenet(config.model)
if checkpoint is not None:
# load model checkpoint
load_checkpoint(model, checkpoint, map_location='cpu')
# save the config in the model for convenience
model.cfg = config
model.to(device)
model.eval()
return model
def _xyxy2xywh(bbox_xyxy):
"""Transform the bbox format from x1y1x2y2 to xywh.
Args:
bbox_xyxy (np.ndarray): Bounding boxes (with scores), shaped (n, 4) or
(n, 5). (left, top, right, bottom, [score])
Returns:
np.ndarray: Bounding boxes (with scores),
shaped (n, 4) or (n, 5). (left, top, width, height, [score])
"""
bbox_xywh = bbox_xyxy.copy()
bbox_xywh[:, 2] = bbox_xywh[:, 2] - bbox_xywh[:, 0] + 1
bbox_xywh[:, 3] = bbox_xywh[:, 3] - bbox_xywh[:, 1] + 1
return bbox_xywh
def _xywh2xyxy(bbox_xywh):
"""Transform the bbox format from xywh to x1y1x2y2.
Args:
bbox_xywh (ndarray): Bounding boxes (with scores),
shaped (n, 4) or (n, 5). (left, top, width, height, [score])
Returns:
np.ndarray: Bounding boxes (with scores), shaped (n, 4) or
(n, 5). (left, top, right, bottom, [score])
"""
bbox_xyxy = bbox_xywh.copy()
bbox_xyxy[:, 2] = bbox_xyxy[:, 2] + bbox_xyxy[:, 0] - 1
bbox_xyxy[:, 3] = bbox_xyxy[:, 3] + bbox_xyxy[:, 1] - 1
return bbox_xyxy
def _box2cs(cfg, box):
"""This encodes bbox(x,y,w,h) into (center, scale)
Args:
x, y, w, h
Returns:
tuple: A tuple containing center and scale.
- np.ndarray[float32](2,): Center of the bbox (x, y).
- np.ndarray[float32](2,): Scale of the bbox w & h.
"""
x, y, w, h = box[:4]
input_size = cfg.data_cfg['image_size']
aspect_ratio = input_size[0] / input_size[1]
center = np.array([x + w * 0.5, y + h * 0.5], dtype=np.float32)
if w > aspect_ratio * h:
h = w * 1.0 / aspect_ratio
elif w < aspect_ratio * h:
w = h * aspect_ratio
# pixel std is 200.0
scale = np.array([w / 200.0, h / 200.0], dtype=np.float32)
scale = scale * 1.25
return center, scale
def _inference_single_pose_model(model,
img_or_path,
bboxes,
dataset='TopDownCocoDataset',
dataset_info=None,
return_heatmap=False):
"""Inference human bounding boxes.
Note:
- num_bboxes: N
- num_keypoints: K
Args:
model (nn.Module): The loaded pose model.
img_or_path (str | np.ndarray): Image filename or loaded image.
bboxes (list | np.ndarray): All bounding boxes (with scores),
shaped (N, 4) or (N, 5). (left, top, width, height, [score])
where N is number of bounding boxes.
dataset (str): Dataset name. Deprecated.
dataset_info (DatasetInfo): A class containing all dataset info.
outputs (list[str] | tuple[str]): Names of layers whose output is
to be returned, default: None
Returns:
ndarray[NxKx3]: Predicted pose x, y, score.
heatmap[N, K, H, W]: Model output heatmap.
"""
cfg = model.cfg
device = next(model.parameters()).device
if device.type == 'cpu':
device = -1
# build the data pipeline
test_pipeline = Compose(cfg.test_pipeline)
assert len(bboxes[0]) in [4, 5]
if dataset_info is not None:
dataset_name = dataset_info.dataset_name
flip_pairs = dataset_info.flip_pairs
else:
warnings.warn(
'dataset is deprecated.'
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
# TODO: These will be removed in the later versions.
if dataset in ('TopDownCocoDataset', 'TopDownOCHumanDataset',
'AnimalMacaqueDataset'):
flip_pairs = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12],
[13, 14], [15, 16]]
elif dataset == 'TopDownCocoWholeBodyDataset':
body = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12],
[13, 14], [15, 16]]
foot = [[17, 20], [18, 21], [19, 22]]
face = [[23, 39], [24, 38], [25, 37], [26, 36], [27, 35], [28, 34],
[29, 33], [30, 32], [40, 49], [41, 48], [42, 47], [43, 46],
[44, 45], [54, 58], [55, 57], [59, 68], [60, 67], [61, 66],
[62, 65], [63, 70], [64, 69], [71, 77], [72, 76], [73, 75],
[78, 82], [79, 81], [83, 87], [84, 86], [88, 90]]
hand = [[91, 112], [92, 113], [93, 114], [94, 115], [95, 116],
[96, 117], [97, 118], [98, 119], [99, 120], [100, 121],
[101, 122], [102, 123], [103, 124], [104, 125], [105, 126],
[106, 127], [107, 128], [108, 129], [109, 130], [110, 131],
[111, 132]]
flip_pairs = body + foot + face + hand
elif dataset == 'TopDownAicDataset':
flip_pairs = [[0, 3], [1, 4], [2, 5], [6, 9], [7, 10], [8, 11]]
elif dataset == 'TopDownMpiiDataset':
flip_pairs = [[0, 5], [1, 4], [2, 3], [10, 15], [11, 14], [12, 13]]
elif dataset == 'TopDownMpiiTrbDataset':
flip_pairs = [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11],
[14, 15], [16, 22], [28, 34], [17, 23], [29, 35],
[18, 24], [30, 36], [19, 25], [31, 37], [20, 26],
[32, 38], [21, 27], [33, 39]]
elif dataset in ('OneHand10KDataset', 'FreiHandDataset',
'PanopticDataset', 'InterHand2DDataset'):
flip_pairs = []
elif dataset in 'Face300WDataset':
flip_pairs = [[0, 16], [1, 15], [2, 14], [3, 13], [4, 12], [5, 11],
[6, 10], [7, 9], [17, 26], [18, 25], [19, 24],
[20, 23], [21, 22], [31, 35], [32, 34], [36, 45],
[37, 44], [38, 43], [39, 42], [40, 47], [41, 46],
[48, 54], [49, 53], [50, 52], [61, 63], [60, 64],
[67, 65], [58, 56], [59, 55]]
elif dataset in 'FaceAFLWDataset':
flip_pairs = [[0, 5], [1, 4], [2, 3], [6, 11], [7, 10], [8, 9],
[12, 14], [15, 17]]
elif dataset in 'FaceCOFWDataset':
flip_pairs = [[0, 1], [4, 6], [2, 3], [5, 7], [8, 9], [10, 11],
[12, 14], [16, 17], [13, 15], [18, 19], [22, 23]]
elif dataset in 'FaceWFLWDataset':
flip_pairs = [[0, 32], [1, 31], [2, 30], [3, 29], [4, 28], [5, 27],
[6, 26], [7, 25], [8, 24], [9, 23], [10, 22],
[11, 21], [12, 20], [13, 19], [14, 18], [15, 17],
[33, 46], [34, 45], [35, 44], [36, 43], [37, 42],
[38, 50], [39, 49], [40, 48], [41, 47], [60, 72],
[61, 71], [62, 70], [63, 69], [64, 68], [65, 75],
[66, 74], [67, 73], [55, 59], [56, 58], [76, 82],
[77, 81], [78, 80], [87, 83], [86, 84], [88, 92],
[89, 91], [95, 93], [96, 97]]
elif dataset in 'AnimalFlyDataset':
flip_pairs = [[1, 2], [6, 18], [7, 19], [8, 20], [9, 21], [10, 22],
[11, 23], [12, 24], [13, 25], [14, 26], [15, 27],
[16, 28], [17, 29], [30, 31]]
elif dataset in 'AnimalHorse10Dataset':
flip_pairs = []
elif dataset in 'AnimalLocustDataset':
flip_pairs = [[5, 20], [6, 21], [7, 22], [8, 23], [9, 24],
[10, 25], [11, 26], [12, 27], [13, 28], [14, 29],
[15, 30], [16, 31], [17, 32], [18, 33], [19, 34]]
elif dataset in 'AnimalZebraDataset':
flip_pairs = [[3, 4], [5, 6]]
elif dataset in 'AnimalPoseDataset':
flip_pairs = [[0, 1], [2, 3], [8, 9], [10, 11], [12, 13], [14, 15],
[16, 17], [18, 19]]
else:
raise NotImplementedError()
dataset_name = dataset
batch_data = []
for bbox in bboxes:
center, scale = _box2cs(cfg, bbox)
# prepare data
data = {
'center':
center,
'scale':
scale,
'bbox_score':
bbox[4] if len(bbox) == 5 else 1,
'bbox_id':
0, # need to be assigned if batch_size > 1
'dataset':
dataset_name,
'joints_3d':
np.zeros((cfg.data_cfg.num_joints, 3), dtype=np.float32),
'joints_3d_visible':
np.zeros((cfg.data_cfg.num_joints, 3), dtype=np.float32),
'rotation':
0,
'ann_info': {
'image_size': np.array(cfg.data_cfg['image_size']),
'num_joints': cfg.data_cfg['num_joints'],
'flip_pairs': flip_pairs
}
}
if isinstance(img_or_path, np.ndarray):
data['img'] = img_or_path
else:
data['image_file'] = img_or_path
data = test_pipeline(data)
batch_data.append(data)
batch_data = collate(batch_data, samples_per_gpu=len(batch_data))
batch_data = scatter(batch_data, [device])[0]
# forward the model
with torch.no_grad():
result = model(
img=batch_data['img'],
img_metas=batch_data['img_metas'],
return_loss=False,
return_heatmap=return_heatmap)
return result['preds'], result['output_heatmap']
def inference_top_down_pose_model(model,
img_or_path,
person_results=None,
bbox_thr=None,
format='xywh',
dataset='TopDownCocoDataset',
dataset_info=None,
return_heatmap=False,
outputs=None):
"""Inference a single image with a list of person bounding boxes.
Note:
- num_people: P
- num_keypoints: K
- bbox height: H
- bbox width: W
Args:
model (nn.Module): The loaded pose model.
img_or_path (str| np.ndarray): Image filename or loaded image.
person_results (list(dict), optional): a list of detected persons that
contains ``bbox`` and/or ``track_id``:
- ``bbox`` (4, ) or (5, ): The person bounding box, which contains
4 box coordinates (and score).
- ``track_id`` (int): The unique id for each human instance. If
not provided, a dummy person result with a bbox covering
the entire image will be used. Default: None.
bbox_thr (float | None): Threshold for bounding boxes. Only bboxes
with higher scores will be fed into the pose detector.
If bbox_thr is None, all boxes will be used.
format (str): bbox format ('xyxy' | 'xywh'). Default: 'xywh'.
- `xyxy` means (left, top, right, bottom),
- `xywh` means (left, top, width, height).
dataset (str): Dataset name, e.g. 'TopDownCocoDataset'.
It is deprecated. Please use dataset_info instead.
dataset_info (DatasetInfo): A class containing all dataset info.
return_heatmap (bool) : Flag to return heatmap, default: False
outputs (list(str) | tuple(str)) : Names of layers whose outputs
need to be returned. Default: None.
Returns:
tuple:
- pose_results (list[dict]): The bbox & pose info. \
Each item in the list is a dictionary, \
containing the bbox: (left, top, right, bottom, [score]) \
and the pose (ndarray[Kx3]): x, y, score.
- returned_outputs (list[dict[np.ndarray[N, K, H, W] | \
torch.Tensor[N, K, H, W]]]): \
Output feature maps from layers specified in `outputs`. \
Includes 'heatmap' if `return_heatmap` is True.
"""
# get dataset info
if (dataset_info is None and hasattr(model, 'cfg')
and 'dataset_info' in model.cfg):
dataset_info = DatasetInfo(model.cfg.dataset_info)
if dataset_info is None:
warnings.warn(
'dataset is deprecated.'
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663'
' for details.', DeprecationWarning)
# only two kinds of bbox format is supported.
assert format in ['xyxy', 'xywh']
pose_results = []
returned_outputs = []
if person_results is None:
# create dummy person results
if isinstance(img_or_path, str):
width, height = Image.open(img_or_path).size
else:
height, width = img_or_path.shape[:2]
person_results = [{'bbox': np.array([0, 0, width, height])}]
if len(person_results) == 0:
return pose_results, returned_outputs
# Change for-loop preprocess each bbox to preprocess all bboxes at once.
bboxes = np.array([box['bbox'] for box in person_results])
# Select bboxes by score threshold
if bbox_thr is not None:
assert bboxes.shape[1] == 5
valid_idx = np.where(bboxes[:, 4] > bbox_thr)[0]
bboxes = bboxes[valid_idx]
person_results = [person_results[i] for i in valid_idx]
if format == 'xyxy':
bboxes_xyxy = bboxes
bboxes_xywh = _xyxy2xywh(bboxes)
else:
# format is already 'xywh'
bboxes_xywh = bboxes
bboxes_xyxy = _xywh2xyxy(bboxes)
# if bbox_thr remove all bounding box
if len(bboxes_xywh) == 0:
return [], []
with OutputHook(model, outputs=outputs, as_tensor=False) as h:
# poses is results['pred'] # N x 17x 3
poses, heatmap = _inference_single_pose_model(
model,
img_or_path,
bboxes_xywh,
dataset=dataset,
dataset_info=dataset_info,
return_heatmap=return_heatmap)
if return_heatmap:
h.layer_outputs['heatmap'] = heatmap
returned_outputs.append(h.layer_outputs)
assert len(poses) == len(person_results), print(
len(poses), len(person_results), len(bboxes_xyxy))
for pose, person_result, bbox_xyxy in zip(poses, person_results,
bboxes_xyxy):
pose_result = person_result.copy()
pose_result['keypoints'] = pose
pose_result['bbox'] = bbox_xyxy
pose_results.append(pose_result)
return pose_results, returned_outputs
def inference_bottom_up_pose_model(model,
img_or_path,
dataset='BottomUpCocoDataset',
dataset_info=None,
pose_nms_thr=0.9,
return_heatmap=False,
outputs=None):
"""Inference a single image with a bottom-up pose model.
Note:
- num_people: P
- num_keypoints: K
- bbox height: H
- bbox width: W
Args:
model (nn.Module): The loaded pose model.
img_or_path (str| np.ndarray): Image filename or loaded image.
dataset (str): Dataset name, e.g. 'BottomUpCocoDataset'.
It is deprecated. Please use dataset_info instead.
dataset_info (DatasetInfo): A class containing all dataset info.
pose_nms_thr (float): retain oks overlap < pose_nms_thr, default: 0.9.
return_heatmap (bool) : Flag to return heatmap, default: False.
outputs (list(str) | tuple(str)) : Names of layers whose outputs
need to be returned, default: None.
Returns:
tuple:
- pose_results (list[np.ndarray]): The predicted pose info. \
The length of the list is the number of people (P). \
Each item in the list is a ndarray, containing each \
person's pose (np.ndarray[Kx3]): x, y, score.
- returned_outputs (list[dict[np.ndarray[N, K, H, W] | \
torch.Tensor[N, K, H, W]]]): \
Output feature maps from layers specified in `outputs`. \
Includes 'heatmap' if `return_heatmap` is True.
"""
# get dataset info
if (dataset_info is None and hasattr(model, 'cfg')
and 'dataset_info' in model.cfg):
dataset_info = DatasetInfo(model.cfg.dataset_info)
if dataset_info is not None:
dataset_name = dataset_info.dataset_name
flip_index = dataset_info.flip_index
sigmas = getattr(dataset_info, 'sigmas', None)
else:
warnings.warn(
'dataset is deprecated.'
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
assert (dataset == 'BottomUpCocoDataset')
dataset_name = dataset
flip_index = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
sigmas = None
pose_results = []
returned_outputs = []
cfg = model.cfg
device = next(model.parameters()).device
if device.type == 'cpu':
device = -1
# build the data pipeline
test_pipeline = Compose(cfg.test_pipeline)
# prepare data
data = {
'dataset': dataset_name,
'ann_info': {
'image_size': np.array(cfg.data_cfg['image_size']),
'num_joints': cfg.data_cfg['num_joints'],
'flip_index': flip_index,
}
}
if isinstance(img_or_path, np.ndarray):
data['img'] = img_or_path
else:
data['image_file'] = img_or_path
data = test_pipeline(data)
data = collate([data], samples_per_gpu=1)
data = scatter(data, [device])[0]
with OutputHook(model, outputs=outputs, as_tensor=False) as h:
# forward the model
with torch.no_grad():
result = model(
img=data['img'],
img_metas=data['img_metas'],
return_loss=False,
return_heatmap=return_heatmap)
if return_heatmap:
h.layer_outputs['heatmap'] = result['output_heatmap']
returned_outputs.append(h.layer_outputs)
for idx, pred in enumerate(result['preds']):
area = (np.max(pred[:, 0]) - np.min(pred[:, 0])) * (
np.max(pred[:, 1]) - np.min(pred[:, 1]))
pose_results.append({
'keypoints': pred[:, :3],
'score': result['scores'][idx],
'area': area,
})
# pose nms
score_per_joint = cfg.model.test_cfg.get('score_per_joint', False)
keep = oks_nms(
pose_results,
pose_nms_thr,
sigmas,
score_per_joint=score_per_joint)
pose_results = [pose_results[_keep] for _keep in keep]
return pose_results, returned_outputs
def vis_pose_result(model,
img,
result,
radius=4,
thickness=1,
kpt_score_thr=0.3,
bbox_color='green',
dataset='TopDownCocoDataset',
dataset_info=None,
show=False,
out_file=None):
"""Visualize the detection results on the image.
Args:
model (nn.Module): The loaded detector.
img (str | np.ndarray): Image filename or loaded image.
result (list[dict]): The results to draw over `img`
(bbox_result, pose_result).
radius (int): Radius of circles.
thickness (int): Thickness of lines.
kpt_score_thr (float): The threshold to visualize the keypoints.
skeleton (list[tuple()]): Default None.
show (bool): Whether to show the image. Default True.
out_file (str|None): The filename of the output visualization image.
"""
# get dataset info
if (dataset_info is None and hasattr(model, 'cfg')
and 'dataset_info' in model.cfg):
dataset_info = DatasetInfo(model.cfg.dataset_info)
if dataset_info is not None:
skeleton = dataset_info.skeleton
pose_kpt_color = dataset_info.pose_kpt_color
pose_link_color = dataset_info.pose_link_color
else:
warnings.warn(
'dataset is deprecated.'
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
# TODO: These will be removed in the later versions.
palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102],
[230, 230, 0], [255, 153, 255], [153, 204, 255],
[255, 102, 255], [255, 51, 255], [102, 178, 255],
[51, 153, 255], [255, 153, 153], [255, 102, 102],
[255, 51, 51], [153, 255, 153], [102, 255, 102],
[51, 255, 51], [0, 255, 0], [0, 0, 255],
[255, 0, 0], [255, 255, 255]])
if dataset in ('TopDownCocoDataset', 'BottomUpCocoDataset',
'TopDownOCHumanDataset', 'AnimalMacaqueDataset'):
# show the results
skeleton = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12],
[5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9],
[8, 10], [1, 2], [0, 1], [0, 2], [1, 3], [2, 4],
[3, 5], [4, 6]]
pose_link_color = palette[[
0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16
]]
pose_kpt_color = palette[[
16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0
]]
elif dataset == 'TopDownCocoWholeBodyDataset':
# show the results
skeleton = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12],
[5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9],
[8, 10], [1, 2], [0, 1], [0, 2],
[1, 3], [2, 4], [3, 5], [4, 6], [15, 17], [15, 18],
[15, 19], [16, 20], [16, 21], [16, 22], [91, 92],
[92, 93], [93, 94], [94, 95], [91, 96], [96, 97],
[97, 98], [98, 99], [91, 100], [100, 101], [101, 102],
[102, 103], [91, 104], [104, 105], [105, 106],
[106, 107], [91, 108], [108, 109], [109, 110],
[110, 111], [112, 113], [113, 114], [114, 115],
[115, 116], [112, 117], [117, 118], [118, 119],
[119, 120], [112, 121], [121, 122], [122, 123],
[123, 124], [112, 125], [125, 126], [126, 127],
[127, 128], [112, 129], [129, 130], [130, 131],
[131, 132]]
pose_link_color = palette[[
0, 0, 0, 0, 7, 7, 7, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 16
] + [16, 16, 16, 16, 16, 16] + [
0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16,
16
] + [
0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16,
16
]]
pose_kpt_color = palette[
[16, 16, 16, 16, 16, 9, 9, 9, 9, 9, 9, 0, 0, 0, 0, 0, 0] +
[0, 0, 0, 0, 0, 0] + [19] * (68 + 42)]
elif dataset == 'TopDownAicDataset':
skeleton = [[2, 1], [1, 0], [0, 13], [13, 3], [3, 4], [4, 5],
[8, 7], [7, 6], [6, 9], [9, 10], [10, 11], [12, 13],
[0, 6], [3, 9]]
pose_link_color = palette[[
9, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 0, 7, 7
]]
pose_kpt_color = palette[[
9, 9, 9, 9, 9, 9, 16, 16, 16, 16, 16, 16, 0, 0
]]
elif dataset == 'TopDownMpiiDataset':
skeleton = [[0, 1], [1, 2], [2, 6], [6, 3], [3, 4], [4, 5], [6, 7],
[7, 8], [8, 9], [8, 12], [12, 11], [11, 10], [8, 13],
[13, 14], [14, 15]]
pose_link_color = palette[[
16, 16, 16, 16, 16, 16, 7, 7, 0, 9, 9, 9, 9, 9, 9
]]
pose_kpt_color = palette[[
16, 16, 16, 16, 16, 16, 7, 7, 0, 0, 9, 9, 9, 9, 9, 9
]]
elif dataset == 'TopDownMpiiTrbDataset':
skeleton = [[12, 13], [13, 0], [13, 1], [0, 2], [1, 3], [2, 4],
[3, 5], [0, 6], [1, 7], [6, 7], [6, 8], [7,
9], [8, 10],
[9, 11], [14, 15], [16, 17], [18, 19], [20, 21],
[22, 23], [24, 25], [26, 27], [28, 29], [30, 31],
[32, 33], [34, 35], [36, 37], [38, 39]]
pose_link_color = palette[[16] * 14 + [19] * 13]
pose_kpt_color = palette[[16] * 14 + [0] * 26]
elif dataset in ('OneHand10KDataset', 'FreiHandDataset',
'PanopticDataset'):
skeleton = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7],
[7, 8], [0, 9], [9, 10], [10, 11], [11, 12], [0, 13],
[13, 14], [14, 15], [15, 16], [0, 17], [17, 18],
[18, 19], [19, 20]]
pose_link_color = palette[[
0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16,
16
]]
pose_kpt_color = palette[[
0, 0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16,
16, 16
]]
elif dataset == 'InterHand2DDataset':
skeleton = [[0, 1], [1, 2], [2, 3], [4, 5], [5, 6], [6, 7], [8, 9],
[9, 10], [10, 11], [12, 13], [13, 14], [14, 15],
[16, 17], [17, 18], [18, 19], [3, 20], [7, 20],
[11, 20], [15, 20], [19, 20]]
pose_link_color = palette[[
0, 0, 0, 4, 4, 4, 8, 8, 8, 12, 12, 12, 16, 16, 16, 0, 4, 8, 12,
16
]]
pose_kpt_color = palette[[
0, 0, 0, 0, 4, 4, 4, 4, 8, 8, 8, 8, 12, 12, 12, 12, 16, 16, 16,
16, 0
]]
elif dataset == 'Face300WDataset':
# show the results
skeleton = []
pose_link_color = palette[[]]
pose_kpt_color = palette[[19] * 68]
kpt_score_thr = 0
elif dataset == 'FaceAFLWDataset':
# show the results
skeleton = []
pose_link_color = palette[[]]
pose_kpt_color = palette[[19] * 19]
kpt_score_thr = 0
elif dataset == 'FaceCOFWDataset':
# show the results
skeleton = []
pose_link_color = palette[[]]
pose_kpt_color = palette[[19] * 29]
kpt_score_thr = 0
elif dataset == 'FaceWFLWDataset':
# show the results
skeleton = []
pose_link_color = palette[[]]
pose_kpt_color = palette[[19] * 98]
kpt_score_thr = 0
elif dataset == 'AnimalHorse10Dataset':
skeleton = [[0, 1], [1, 12], [12, 16], [16, 21], [21, 17],
[17, 11], [11, 10], [10, 8], [8, 9], [9, 12], [2, 3],
[3, 4], [5, 6], [6, 7], [13, 14], [14, 15], [18, 19],
[19, 20]]
pose_link_color = palette[[4] * 10 + [6] * 2 + [6] * 2 + [7] * 2 +
[7] * 2]
pose_kpt_color = palette[[
4, 4, 6, 6, 6, 6, 6, 6, 4, 4, 4, 4, 4, 7, 7, 7, 4, 4, 7, 7, 7,
4
]]
elif dataset == 'AnimalFlyDataset':
skeleton = [[1, 0], [2, 0], [3, 0], [4, 3], [5, 4], [7, 6], [8, 7],
[9, 8], [11, 10], [12, 11], [13, 12], [15, 14],
[16, 15], [17, 16], [19, 18], [20, 19], [21, 20],
[23, 22], [24, 23], [25, 24], [27, 26], [28, 27],
[29, 28], [30, 3], [31, 3]]
pose_link_color = palette[[0] * 25]
pose_kpt_color = palette[[0] * 32]
elif dataset == 'AnimalLocustDataset':
skeleton = [[1, 0], [2, 1], [3, 2], [4, 3], [6, 5], [7, 6], [9, 8],
[10, 9], [11, 10], [13, 12], [14, 13], [15, 14],
[17, 16], [18, 17], [19, 18], [21, 20], [22, 21],
[24, 23], [25, 24], [26, 25], [28, 27], [29, 28],
[30, 29], [32, 31], [33, 32], [34, 33]]
pose_link_color = palette[[0] * 26]
pose_kpt_color = palette[[0] * 35]
elif dataset == 'AnimalZebraDataset':
skeleton = [[1, 0], [2, 1], [3, 2], [4, 2], [5, 7], [6, 7], [7, 2],
[8, 7]]
pose_link_color = palette[[0] * 8]
pose_kpt_color = palette[[0] * 9]
elif dataset in 'AnimalPoseDataset':
skeleton = [[0, 1], [0, 2], [1, 3], [0, 4], [1, 4], [4, 5], [5, 7],
[6, 7], [5, 8], [8, 12], [12, 16], [5, 9], [9, 13],
[13, 17], [6, 10], [10, 14], [14, 18], [6, 11],
[11, 15], [15, 19]]
pose_link_color = palette[[0] * 20]
pose_kpt_color = palette[[0] * 20]
else:
NotImplementedError()
if hasattr(model, 'module'):
model = model.module
img = model.show_result(
img,
result,
skeleton,
radius=radius,
thickness=thickness,
pose_kpt_color=pose_kpt_color,
pose_link_color=pose_link_color,
kpt_score_thr=kpt_score_thr,
bbox_color=bbox_color,
show=show,
out_file=out_file)
return img
def process_mmdet_results(mmdet_results, cat_id=1):
"""Process mmdet results, and return a list of bboxes.
Args:
mmdet_results (list|tuple): mmdet results.
cat_id (int): category id (default: 1 for human)
Returns:
person_results (list): a list of detected bounding boxes
"""
if isinstance(mmdet_results, tuple):
det_results = mmdet_results[0]
else:
det_results = mmdet_results
bboxes = det_results[cat_id - 1]
person_results = []
for bbox in bboxes:
person = {}
person['bbox'] = bbox
person_results.append(person)
return person_results
third_party/hamer/third-party/ViTPose/mmpose/apis/inference_3d.py
Vendored
+791
View File
@@ -0,0 +1,791 @@
# Copyright (c) OpenMMLab. All rights reserved.
import warnings
import numpy as np
import torch
from mmcv.parallel import collate, scatter
from mmpose.datasets.pipelines import Compose
from .inference import _box2cs, _xywh2xyxy, _xyxy2xywh
def extract_pose_sequence(pose_results, frame_idx, causal, seq_len, step=1):
"""Extract the target frame from 2D pose results, and pad the sequence to a
fixed length.
Args:
pose_results (list[list[dict]]): Multi-frame pose detection results
stored in a nested list. Each element of the outer list is the
pose detection results of a single frame, and each element of the
inner list is the pose information of one person, which contains:
- keypoints (ndarray[K, 2 or 3]): x, y, [score]
- track_id (int): unique id of each person, required \
when ``with_track_id==True``.
- bbox ((4, ) or (5, )): left, right, top, bottom, [score]
frame_idx (int): The index of the frame in the original video.
causal (bool): If True, the target frame is the last frame in
a sequence. Otherwise, the target frame is in the middle of
a sequence.
seq_len (int): The number of frames in the input sequence.
step (int): Step size to extract frames from the video.
Returns:
list[list[dict]]: Multi-frame pose detection results stored \
in a nested list with a length of seq_len.
"""
if causal:
frames_left = seq_len - 1
frames_right = 0
else:
frames_left = (seq_len - 1) // 2
frames_right = frames_left
num_frames = len(pose_results)
# get the padded sequence
pad_left = max(0, frames_left - frame_idx // step)
pad_right = max(0, frames_right - (num_frames - 1 - frame_idx) // step)
start = max(frame_idx % step, frame_idx - frames_left * step)
end = min(num_frames - (num_frames - 1 - frame_idx) % step,
frame_idx + frames_right * step + 1)
pose_results_seq = [pose_results[0]] * pad_left + \
pose_results[start:end:step] + [pose_results[-1]] * pad_right
return pose_results_seq
def _gather_pose_lifter_inputs(pose_results,
bbox_center,
bbox_scale,
norm_pose_2d=False):
"""Gather input data (keypoints and track_id) for pose lifter model.
Note:
- The temporal length of the pose detection results: T
- The number of the person instances: N
- The number of the keypoints: K
- The channel number of each keypoint: C
Args:
pose_results (List[List[Dict]]): Multi-frame pose detection results
stored in a nested list. Each element of the outer list is the
pose detection results of a single frame, and each element of the
inner list is the pose information of one person, which contains:
- keypoints (ndarray[K, 2 or 3]): x, y, [score]
- track_id (int): unique id of each person, required when
``with_track_id==True```
- bbox ((4, ) or (5, )): left, right, top, bottom, [score]
bbox_center (ndarray[1, 2]): x, y. The average center coordinate of the
bboxes in the dataset.
bbox_scale (int|float): The average scale of the bboxes in the dataset.
norm_pose_2d (bool): If True, scale the bbox (along with the 2D
pose) to bbox_scale, and move the bbox (along with the 2D pose) to
bbox_center. Default: False.
Returns:
list[list[dict]]: Multi-frame pose detection results
stored in a nested list. Each element of the outer list is the
pose detection results of a single frame, and each element of the
inner list is the pose information of one person, which contains:
- keypoints (ndarray[K, 2 or 3]): x, y, [score]
- track_id (int): unique id of each person, required when
``with_track_id==True``
"""
sequence_inputs = []
for frame in pose_results:
frame_inputs = []
for res in frame:
inputs = dict()
if norm_pose_2d:
bbox = res['bbox']
center = np.array([[(bbox[0] + bbox[2]) / 2,
(bbox[1] + bbox[3]) / 2]])
scale = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
inputs['keypoints'] = (res['keypoints'][:, :2] - center) \
/ scale * bbox_scale + bbox_center
else:
inputs['keypoints'] = res['keypoints'][:, :2]
if res['keypoints'].shape[1] == 3:
inputs['keypoints'] = np.concatenate(
[inputs['keypoints'], res['keypoints'][:, 2:]], axis=1)
if 'track_id' in res:
inputs['track_id'] = res['track_id']
frame_inputs.append(inputs)
sequence_inputs.append(frame_inputs)
return sequence_inputs
def _collate_pose_sequence(pose_results, with_track_id=True, target_frame=-1):
"""Reorganize multi-frame pose detection results into individual pose
sequences.
Note:
- The temporal length of the pose detection results: T
- The number of the person instances: N
- The number of the keypoints: K
- The channel number of each keypoint: C
Args:
pose_results (List[List[Dict]]): Multi-frame pose detection results
stored in a nested list. Each element of the outer list is the
pose detection results of a single frame, and each element of the
inner list is the pose information of one person, which contains:
- keypoints (ndarray[K, 2 or 3]): x, y, [score]
- track_id (int): unique id of each person, required when
``with_track_id==True```
with_track_id (bool): If True, the element in pose_results is expected
to contain "track_id", which will be used to gather the pose
sequence of a person from multiple frames. Otherwise, the pose
results in each frame are expected to have a consistent number and
order of identities. Default is True.
target_frame (int): The index of the target frame. Default: -1.
"""
T = len(pose_results)
assert T > 0
target_frame = (T + target_frame) % T # convert negative index to positive
N = len(pose_results[target_frame]) # use identities in the target frame
if N == 0:
return []
K, C = pose_results[target_frame][0]['keypoints'].shape
track_ids = None
if with_track_id:
track_ids = [res['track_id'] for res in pose_results[target_frame]]
pose_sequences = []
for idx in range(N):
pose_seq = dict()
# gather static information
for k, v in pose_results[target_frame][idx].items():
if k != 'keypoints':
pose_seq[k] = v
# gather keypoints
if not with_track_id:
pose_seq['keypoints'] = np.stack(
[frame[idx]['keypoints'] for frame in pose_results])
else:
keypoints = np.zeros((T, K, C), dtype=np.float32)
keypoints[target_frame] = pose_results[target_frame][idx][
'keypoints']
# find the left most frame containing track_ids[idx]
for frame_idx in range(target_frame - 1, -1, -1):
contains_idx = False
for res in pose_results[frame_idx]:
if res['track_id'] == track_ids[idx]:
keypoints[frame_idx] = res['keypoints']
contains_idx = True
break
if not contains_idx:
# replicate the left most frame
keypoints[:frame_idx + 1] = keypoints[frame_idx + 1]
break
# find the right most frame containing track_idx[idx]
for frame_idx in range(target_frame + 1, T):
contains_idx = False
for res in pose_results[frame_idx]:
if res['track_id'] == track_ids[idx]:
keypoints[frame_idx] = res['keypoints']
contains_idx = True
break
if not contains_idx:
# replicate the right most frame
keypoints[frame_idx + 1:] = keypoints[frame_idx]
break
pose_seq['keypoints'] = keypoints
pose_sequences.append(pose_seq)
return pose_sequences
def inference_pose_lifter_model(model,
pose_results_2d,
dataset=None,
dataset_info=None,
with_track_id=True,
image_size=None,
norm_pose_2d=False):
"""Inference 3D pose from 2D pose sequences using a pose lifter model.
Args:
model (nn.Module): The loaded pose lifter model
pose_results_2d (list[list[dict]]): The 2D pose sequences stored in a
nested list. Each element of the outer list is the 2D pose results
of a single frame, and each element of the inner list is the 2D
pose of one person, which contains:
- "keypoints" (ndarray[K, 2 or 3]): x, y, [score]
- "track_id" (int)
dataset (str): Dataset name, e.g. 'Body3DH36MDataset'
with_track_id: If True, the element in pose_results_2d is expected to
contain "track_id", which will be used to gather the pose sequence
of a person from multiple frames. Otherwise, the pose results in
each frame are expected to have a consistent number and order of
identities. Default is True.
image_size (tuple|list): image width, image height. If None, image size
will not be contained in dict ``data``.
norm_pose_2d (bool): If True, scale the bbox (along with the 2D
pose) to the average bbox scale of the dataset, and move the bbox
(along with the 2D pose) to the average bbox center of the dataset.
Returns:
list[dict]: 3D pose inference results. Each element is the result of \
an instance, which contains:
- "keypoints_3d" (ndarray[K, 3]): predicted 3D keypoints
- "keypoints" (ndarray[K, 2 or 3]): from the last frame in \
``pose_results_2d``.
- "track_id" (int): from the last frame in ``pose_results_2d``. \
If there is no valid instance, an empty list will be \
returned.
"""
cfg = model.cfg
test_pipeline = Compose(cfg.test_pipeline)
device = next(model.parameters()).device
if device.type == 'cpu':
device = -1
if dataset_info is not None:
flip_pairs = dataset_info.flip_pairs
assert 'stats_info' in dataset_info._dataset_info
bbox_center = dataset_info._dataset_info['stats_info']['bbox_center']
bbox_scale = dataset_info._dataset_info['stats_info']['bbox_scale']
else:
warnings.warn(
'dataset is deprecated.'
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
# TODO: These will be removed in the later versions.
if dataset == 'Body3DH36MDataset':
flip_pairs = [[1, 4], [2, 5], [3, 6], [11, 14], [12, 15], [13, 16]]
bbox_center = np.array([[528, 427]], dtype=np.float32)
bbox_scale = 400
else:
raise NotImplementedError()
target_idx = -1 if model.causal else len(pose_results_2d) // 2
pose_lifter_inputs = _gather_pose_lifter_inputs(pose_results_2d,
bbox_center, bbox_scale,
norm_pose_2d)
pose_sequences_2d = _collate_pose_sequence(pose_lifter_inputs,
with_track_id, target_idx)
if not pose_sequences_2d:
return []
batch_data = []
for seq in pose_sequences_2d:
pose_2d = seq['keypoints'].astype(np.float32)
T, K, C = pose_2d.shape
input_2d = pose_2d[..., :2]
input_2d_visible = pose_2d[..., 2:3]
if C > 2:
input_2d_visible = pose_2d[..., 2:3]
else:
input_2d_visible = np.ones((T, K, 1), dtype=np.float32)
# TODO: Will be removed in the later versions
# Dummy 3D input
# This is for compatibility with configs in mmpose<=v0.14.0, where a
# 3D input is required to generate denormalization parameters. This
# part will be removed in the future.
target = np.zeros((K, 3), dtype=np.float32)
target_visible = np.ones((K, 1), dtype=np.float32)
# Dummy image path
# This is for compatibility with configs in mmpose<=v0.14.0, where
# target_image_path is required. This part will be removed in the
# future.
target_image_path = None
data = {
'input_2d': input_2d,
'input_2d_visible': input_2d_visible,
'target': target,
'target_visible': target_visible,
'target_image_path': target_image_path,
'ann_info': {
'num_joints': K,
'flip_pairs': flip_pairs
}
}
if image_size is not None:
assert len(image_size) == 2
data['image_width'] = image_size[0]
data['image_height'] = image_size[1]
data = test_pipeline(data)
batch_data.append(data)
batch_data = collate(batch_data, samples_per_gpu=len(batch_data))
batch_data = scatter(batch_data, target_gpus=[device])[0]
with torch.no_grad():
result = model(
input=batch_data['input'],
metas=batch_data['metas'],
return_loss=False)
poses_3d = result['preds']
if poses_3d.shape[-1] != 4:
assert poses_3d.shape[-1] == 3
dummy_score = np.ones(
poses_3d.shape[:-1] + (1, ), dtype=poses_3d.dtype)
poses_3d = np.concatenate((poses_3d, dummy_score), axis=-1)
pose_results = []
for pose_2d, pose_3d in zip(pose_sequences_2d, poses_3d):
pose_result = pose_2d.copy()
pose_result['keypoints_3d'] = pose_3d
pose_results.append(pose_result)
return pose_results
def vis_3d_pose_result(model,
result,
img=None,
dataset='Body3DH36MDataset',
dataset_info=None,
kpt_score_thr=0.3,
radius=8,
thickness=2,
num_instances=-1,
show=False,
out_file=None):
"""Visualize the 3D pose estimation results.
Args:
model (nn.Module): The loaded model.
result (list[dict])
"""
if dataset_info is not None:
skeleton = dataset_info.skeleton
pose_kpt_color = dataset_info.pose_kpt_color
pose_link_color = dataset_info.pose_link_color
else:
warnings.warn(
'dataset is deprecated.'
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
# TODO: These will be removed in the later versions.
palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102],
[230, 230, 0], [255, 153, 255], [153, 204, 255],
[255, 102, 255], [255, 51, 255], [102, 178, 255],
[51, 153, 255], [255, 153, 153], [255, 102, 102],
[255, 51, 51], [153, 255, 153], [102, 255, 102],
[51, 255, 51], [0, 255, 0], [0, 0, 255],
[255, 0, 0], [255, 255, 255]])
if dataset == 'Body3DH36MDataset':
skeleton = [[0, 1], [1, 2], [2, 3], [0, 4], [4, 5], [5, 6], [0, 7],
[7, 8], [8, 9], [9, 10], [8, 11], [11, 12], [12, 13],
[8, 14], [14, 15], [15, 16]]
pose_kpt_color = palette[[
9, 0, 0, 0, 16, 16, 16, 9, 9, 9, 9, 16, 16, 16, 0, 0, 0
]]
pose_link_color = palette[[
0, 0, 0, 16, 16, 16, 9, 9, 9, 9, 16, 16, 16, 0, 0, 0
]]
elif dataset == 'InterHand3DDataset':
skeleton = [[0, 1], [1, 2], [2, 3], [3, 20], [4, 5], [5, 6],
[6, 7], [7, 20], [8, 9], [9, 10], [10, 11], [11, 20],
[12, 13], [13, 14], [14, 15], [15, 20], [16, 17],
[17, 18], [18, 19], [19, 20], [21, 22], [22, 23],
[23, 24], [24, 41], [25, 26], [26, 27], [27, 28],
[28, 41], [29, 30], [30, 31], [31, 32], [32, 41],
[33, 34], [34, 35], [35, 36], [36, 41], [37, 38],
[38, 39], [39, 40], [40, 41]]
pose_kpt_color = [[14, 128, 250], [14, 128, 250], [14, 128, 250],
[14, 128, 250], [80, 127, 255], [80, 127, 255],
[80, 127, 255], [80, 127, 255], [71, 99, 255],
[71, 99, 255], [71, 99, 255], [71, 99, 255],
[0, 36, 255], [0, 36, 255], [0, 36, 255],
[0, 36, 255], [0, 0, 230], [0, 0, 230],
[0, 0, 230], [0, 0, 230], [0, 0, 139],
[237, 149, 100], [237, 149, 100],
[237, 149, 100], [237, 149, 100], [230, 128, 77],
[230, 128, 77], [230, 128, 77], [230, 128, 77],
[255, 144, 30], [255, 144, 30], [255, 144, 30],
[255, 144, 30], [153, 51, 0], [153, 51, 0],
[153, 51, 0], [153, 51, 0], [255, 51, 13],
[255, 51, 13], [255, 51, 13], [255, 51, 13],
[103, 37, 8]]
pose_link_color = [[14, 128, 250], [14, 128, 250], [14, 128, 250],
[14, 128, 250], [80, 127, 255], [80, 127, 255],
[80, 127, 255], [80, 127, 255], [71, 99, 255],
[71, 99, 255], [71, 99, 255], [71, 99, 255],
[0, 36, 255], [0, 36, 255], [0, 36, 255],
[0, 36, 255], [0, 0, 230], [0, 0, 230],
[0, 0, 230], [0, 0, 230], [237, 149, 100],
[237, 149, 100], [237, 149, 100],
[237, 149, 100], [230, 128, 77], [230, 128, 77],
[230, 128, 77], [230, 128, 77], [255, 144, 30],
[255, 144, 30], [255, 144, 30], [255, 144, 30],
[153, 51, 0], [153, 51, 0], [153, 51, 0],
[153, 51, 0], [255, 51, 13], [255, 51, 13],
[255, 51, 13], [255, 51, 13]]
else:
raise NotImplementedError
if hasattr(model, 'module'):
model = model.module
img = model.show_result(
result,
img,
skeleton,
radius=radius,
thickness=thickness,
pose_kpt_color=pose_kpt_color,
pose_link_color=pose_link_color,
num_instances=num_instances,
show=show,
out_file=out_file)
return img
def inference_interhand_3d_model(model,
img_or_path,
det_results,
bbox_thr=None,
format='xywh',
dataset='InterHand3DDataset'):
"""Inference a single image with a list of hand bounding boxes.
Note:
- num_bboxes: N
- num_keypoints: K
Args:
model (nn.Module): The loaded pose model.
img_or_path (str | np.ndarray): Image filename or loaded image.
det_results (list[dict]): The 2D bbox sequences stored in a list.
Each each element of the list is the bbox of one person, whose
shape is (ndarray[4 or 5]), containing 4 box coordinates
(and score).
dataset (str): Dataset name.
format: bbox format ('xyxy' | 'xywh'). Default: 'xywh'.
'xyxy' means (left, top, right, bottom),
'xywh' means (left, top, width, height).
Returns:
list[dict]: 3D pose inference results. Each element is the result \
of an instance, which contains the predicted 3D keypoints with \
shape (ndarray[K,3]). If there is no valid instance, an \
empty list will be returned.
"""
assert format in ['xyxy', 'xywh']
pose_results = []
if len(det_results) == 0:
return pose_results
# Change for-loop preprocess each bbox to preprocess all bboxes at once.
bboxes = np.array([box['bbox'] for box in det_results])
# Select bboxes by score threshold
if bbox_thr is not None:
assert bboxes.shape[1] == 5
valid_idx = np.where(bboxes[:, 4] > bbox_thr)[0]
bboxes = bboxes[valid_idx]
det_results = [det_results[i] for i in valid_idx]
if format == 'xyxy':
bboxes_xyxy = bboxes
bboxes_xywh = _xyxy2xywh(bboxes)
else:
# format is already 'xywh'
bboxes_xywh = bboxes
bboxes_xyxy = _xywh2xyxy(bboxes)
# if bbox_thr remove all bounding box
if len(bboxes_xywh) == 0:
return []
cfg = model.cfg
device = next(model.parameters()).device
if device.type == 'cpu':
device = -1
# build the data pipeline
test_pipeline = Compose(cfg.test_pipeline)
assert len(bboxes[0]) in [4, 5]
if dataset == 'InterHand3DDataset':
flip_pairs = [[i, 21 + i] for i in range(21)]
else:
raise NotImplementedError()
batch_data = []
for bbox in bboxes:
center, scale = _box2cs(cfg, bbox)
# prepare data
data = {
'center':
center,
'scale':
scale,
'bbox_score':
bbox[4] if len(bbox) == 5 else 1,
'bbox_id':
0, # need to be assigned if batch_size > 1
'dataset':
dataset,
'joints_3d':
np.zeros((cfg.data_cfg.num_joints, 3), dtype=np.float32),
'joints_3d_visible':
np.zeros((cfg.data_cfg.num_joints, 3), dtype=np.float32),
'rotation':
0,
'ann_info': {
'image_size': np.array(cfg.data_cfg['image_size']),
'num_joints': cfg.data_cfg['num_joints'],
'flip_pairs': flip_pairs,
'heatmap3d_depth_bound': cfg.data_cfg['heatmap3d_depth_bound'],
'heatmap_size_root': cfg.data_cfg['heatmap_size_root'],
'root_depth_bound': cfg.data_cfg['root_depth_bound']
}
}
if isinstance(img_or_path, np.ndarray):
data['img'] = img_or_path
else:
data['image_file'] = img_or_path
data = test_pipeline(data)
batch_data.append(data)
batch_data = collate(batch_data, samples_per_gpu=len(batch_data))
batch_data = scatter(batch_data, [device])[0]
# forward the model
with torch.no_grad():
result = model(
img=batch_data['img'],
img_metas=batch_data['img_metas'],
return_loss=False)
poses_3d = result['preds']
rel_root_depth = result['rel_root_depth']
hand_type = result['hand_type']
if poses_3d.shape[-1] != 4:
assert poses_3d.shape[-1] == 3
dummy_score = np.ones(
poses_3d.shape[:-1] + (1, ), dtype=poses_3d.dtype)
poses_3d = np.concatenate((poses_3d, dummy_score), axis=-1)
# add relative root depth to left hand joints
poses_3d[:, 21:, 2] += rel_root_depth
# set joint scores according to hand type
poses_3d[:, :21, 3] *= hand_type[:, [0]]
poses_3d[:, 21:, 3] *= hand_type[:, [1]]
pose_results = []
for pose_3d, person_res, bbox_xyxy in zip(poses_3d, det_results,
bboxes_xyxy):
pose_res = person_res.copy()
pose_res['keypoints_3d'] = pose_3d
pose_res['bbox'] = bbox_xyxy
pose_results.append(pose_res)
return pose_results
def inference_mesh_model(model,
img_or_path,
det_results,
bbox_thr=None,
format='xywh',
dataset='MeshH36MDataset'):
"""Inference a single image with a list of bounding boxes.
Note:
- num_bboxes: N
- num_keypoints: K
- num_vertices: V
- num_faces: F
Args:
model (nn.Module): The loaded pose model.
img_or_path (str | np.ndarray): Image filename or loaded image.
det_results (list[dict]): The 2D bbox sequences stored in a list.
Each element of the list is the bbox of one person.
"bbox" (ndarray[4 or 5]): The person bounding box,
which contains 4 box coordinates (and score).
bbox_thr (float | None): Threshold for bounding boxes.
Only bboxes with higher scores will be fed into the pose
detector. If bbox_thr is None, all boxes will be used.
format (str): bbox format ('xyxy' | 'xywh'). Default: 'xywh'.
- 'xyxy' means (left, top, right, bottom),
- 'xywh' means (left, top, width, height).
dataset (str): Dataset name.
Returns:
list[dict]: 3D pose inference results. Each element \
is the result of an instance, which contains:
- 'bbox' (ndarray[4]): instance bounding bbox
- 'center' (ndarray[2]): bbox center
- 'scale' (ndarray[2]): bbox scale
- 'keypoints_3d' (ndarray[K,3]): predicted 3D keypoints
- 'camera' (ndarray[3]): camera parameters
- 'vertices' (ndarray[V, 3]): predicted 3D vertices
- 'faces' (ndarray[F, 3]): mesh faces
If there is no valid instance, an empty list
will be returned.
"""
assert format in ['xyxy', 'xywh']
pose_results = []
if len(det_results) == 0:
return pose_results
# Change for-loop preprocess each bbox to preprocess all bboxes at once.
bboxes = np.array([box['bbox'] for box in det_results])
# Select bboxes by score threshold
if bbox_thr is not None:
assert bboxes.shape[1] == 5
valid_idx = np.where(bboxes[:, 4] > bbox_thr)[0]
bboxes = bboxes[valid_idx]
det_results = [det_results[i] for i in valid_idx]
if format == 'xyxy':
bboxes_xyxy = bboxes
bboxes_xywh = _xyxy2xywh(bboxes)
else:
# format is already 'xywh'
bboxes_xywh = bboxes
bboxes_xyxy = _xywh2xyxy(bboxes)
# if bbox_thr remove all bounding box
if len(bboxes_xywh) == 0:
return []
cfg = model.cfg
device = next(model.parameters()).device
if device.type == 'cpu':
device = -1
# build the data pipeline
test_pipeline = Compose(cfg.test_pipeline)
assert len(bboxes[0]) in [4, 5]
if dataset == 'MeshH36MDataset':
flip_pairs = [[0, 5], [1, 4], [2, 3], [6, 11], [7, 10], [8, 9],
[20, 21], [22, 23]]
else:
raise NotImplementedError()
batch_data = []
for bbox in bboxes:
center, scale = _box2cs(cfg, bbox)
# prepare data
data = {
'image_file':
img_or_path,
'center':
center,
'scale':
scale,
'rotation':
0,
'bbox_score':
bbox[4] if len(bbox) == 5 else 1,
'dataset':
dataset,
'joints_2d':
np.zeros((cfg.data_cfg.num_joints, 2), dtype=np.float32),
'joints_2d_visible':
np.zeros((cfg.data_cfg.num_joints, 1), dtype=np.float32),
'joints_3d':
np.zeros((cfg.data_cfg.num_joints, 3), dtype=np.float32),
'joints_3d_visible':
np.zeros((cfg.data_cfg.num_joints, 3), dtype=np.float32),
'pose':
np.zeros(72, dtype=np.float32),
'beta':
np.zeros(10, dtype=np.float32),
'has_smpl':
0,
'ann_info': {
'image_size': np.array(cfg.data_cfg['image_size']),
'num_joints': cfg.data_cfg['num_joints'],
'flip_pairs': flip_pairs,
}
}
data = test_pipeline(data)
batch_data.append(data)
batch_data = collate(batch_data, samples_per_gpu=len(batch_data))
batch_data = scatter(batch_data, target_gpus=[device])[0]
# forward the model
with torch.no_grad():
preds = model(
img=batch_data['img'],
img_metas=batch_data['img_metas'],
return_loss=False,
return_vertices=True,
return_faces=True)
for idx in range(len(det_results)):
pose_res = det_results[idx].copy()
pose_res['bbox'] = bboxes_xyxy[idx]
pose_res['center'] = batch_data['img_metas'][idx]['center']
pose_res['scale'] = batch_data['img_metas'][idx]['scale']
pose_res['keypoints_3d'] = preds['keypoints_3d'][idx]
pose_res['camera'] = preds['camera'][idx]
pose_res['vertices'] = preds['vertices'][idx]
pose_res['faces'] = preds['faces']
pose_results.append(pose_res)
return pose_results
def vis_3d_mesh_result(model, result, img=None, show=False, out_file=None):
"""Visualize the 3D mesh estimation results.
Args:
model (nn.Module): The loaded model.
result (list[dict]): 3D mesh estimation results.
"""
if hasattr(model, 'module'):
model = model.module
img = model.show_result(result, img, show=show, out_file=out_file)
return img
third_party/hamer/third-party/ViTPose/mmpose/apis/inference_tracking.py
Vendored
+347
View File
@@ -0,0 +1,347 @@
# Copyright (c) OpenMMLab. All rights reserved.
import warnings
import numpy as np
from mmpose.core import OneEuroFilter, oks_iou
def _compute_iou(bboxA, bboxB):
"""Compute the Intersection over Union (IoU) between two boxes .
Args:
bboxA (list): The first bbox info (left, top, right, bottom, score).
bboxB (list): The second bbox info (left, top, right, bottom, score).
Returns:
float: The IoU value.
"""
x1 = max(bboxA[0], bboxB[0])
y1 = max(bboxA[1], bboxB[1])
x2 = min(bboxA[2], bboxB[2])
y2 = min(bboxA[3], bboxB[3])
inter_area = max(0, x2 - x1) * max(0, y2 - y1)
bboxA_area = (bboxA[2] - bboxA[0]) * (bboxA[3] - bboxA[1])
bboxB_area = (bboxB[2] - bboxB[0]) * (bboxB[3] - bboxB[1])
union_area = float(bboxA_area + bboxB_area - inter_area)
if union_area == 0:
union_area = 1e-5
warnings.warn('union_area=0 is unexpected')
iou = inter_area / union_area
return iou
def _track_by_iou(res, results_last, thr):
"""Get track id using IoU tracking greedily.
Args:
res (dict): The bbox & pose results of the person instance.
results_last (list[dict]): The bbox & pose & track_id info of the
last frame (bbox_result, pose_result, track_id).
thr (float): The threshold for iou tracking.
Returns:
int: The track id for the new person instance.
list[dict]: The bbox & pose & track_id info of the persons
that have not been matched on the last frame.
dict: The matched person instance on the last frame.
"""
bbox = list(res['bbox'])
max_iou_score = -1
max_index = -1
match_result = {}
for index, res_last in enumerate(results_last):
bbox_last = list(res_last['bbox'])
iou_score = _compute_iou(bbox, bbox_last)
if iou_score > max_iou_score:
max_iou_score = iou_score
max_index = index
if max_iou_score > thr:
track_id = results_last[max_index]['track_id']
match_result = results_last[max_index]
del results_last[max_index]
else:
track_id = -1
return track_id, results_last, match_result
def _track_by_oks(res, results_last, thr):
"""Get track id using OKS tracking greedily.
Args:
res (dict): The pose results of the person instance.
results_last (list[dict]): The pose & track_id info of the
last frame (pose_result, track_id).
thr (float): The threshold for oks tracking.
Returns:
int: The track id for the new person instance.
list[dict]: The pose & track_id info of the persons
that have not been matched on the last frame.
dict: The matched person instance on the last frame.
"""
pose = res['keypoints'].reshape((-1))
area = res['area']
max_index = -1
match_result = {}
if len(results_last) == 0:
return -1, results_last, match_result
pose_last = np.array(
[res_last['keypoints'].reshape((-1)) for res_last in results_last])
area_last = np.array([res_last['area'] for res_last in results_last])
oks_score = oks_iou(pose, pose_last, area, area_last)
max_index = np.argmax(oks_score)
if oks_score[max_index] > thr:
track_id = results_last[max_index]['track_id']
match_result = results_last[max_index]
del results_last[max_index]
else:
track_id = -1
return track_id, results_last, match_result
def _get_area(results):
"""Get bbox for each person instance on the current frame.
Args:
results (list[dict]): The pose results of the current frame
(pose_result).
Returns:
list[dict]: The bbox & pose info of the current frame
(bbox_result, pose_result, area).
"""
for result in results:
if 'bbox' in result:
result['area'] = ((result['bbox'][2] - result['bbox'][0]) *
(result['bbox'][3] - result['bbox'][1]))
else:
xmin = np.min(
result['keypoints'][:, 0][result['keypoints'][:, 0] > 0],
initial=1e10)
xmax = np.max(result['keypoints'][:, 0])
ymin = np.min(
result['keypoints'][:, 1][result['keypoints'][:, 1] > 0],
initial=1e10)
ymax = np.max(result['keypoints'][:, 1])
result['area'] = (xmax - xmin) * (ymax - ymin)
result['bbox'] = np.array([xmin, ymin, xmax, ymax])
return results
def _temporal_refine(result, match_result, fps=None):
"""Refine koypoints using tracked person instance on last frame.
Args:
results (dict): The pose results of the current frame
(pose_result).
match_result (dict): The pose results of the last frame
(match_result)
Returns:
(array): The person keypoints after refine.
"""
if 'one_euro' in match_result:
result['keypoints'][:, :2] = match_result['one_euro'](
result['keypoints'][:, :2])
result['one_euro'] = match_result['one_euro']
else:
result['one_euro'] = OneEuroFilter(result['keypoints'][:, :2], fps=fps)
return result['keypoints']
def get_track_id(results,
results_last,
next_id,
min_keypoints=3,
use_oks=False,
tracking_thr=0.3,
use_one_euro=False,
fps=None):
"""Get track id for each person instance on the current frame.
Args:
results (list[dict]): The bbox & pose results of the current frame
(bbox_result, pose_result).
results_last (list[dict]): The bbox & pose & track_id info of the
last frame (bbox_result, pose_result, track_id).
next_id (int): The track id for the new person instance.
min_keypoints (int): Minimum number of keypoints recognized as person.
default: 3.
use_oks (bool): Flag to using oks tracking. default: False.
tracking_thr (float): The threshold for tracking.
use_one_euro (bool): Option to use one-euro-filter. default: False.
fps (optional): Parameters that d_cutoff
when one-euro-filter is used as a video input
Returns:
tuple:
- results (list[dict]): The bbox & pose & track_id info of the \
current frame (bbox_result, pose_result, track_id).
- next_id (int): The track id for the new person instance.
"""
results = _get_area(results)
if use_oks:
_track = _track_by_oks
else:
_track = _track_by_iou
for result in results:
track_id, results_last, match_result = _track(result, results_last,
tracking_thr)
if track_id == -1:
if np.count_nonzero(result['keypoints'][:, 1]) > min_keypoints:
result['track_id'] = next_id
next_id += 1
else:
# If the number of keypoints detected is small,
# delete that person instance.
result['keypoints'][:, 1] = -10
result['bbox'] *= 0
result['track_id'] = -1
else:
result['track_id'] = track_id
if use_one_euro:
result['keypoints'] = _temporal_refine(
result, match_result, fps=fps)
del match_result
return results, next_id
def vis_pose_tracking_result(model,
img,
result,
radius=4,
thickness=1,
kpt_score_thr=0.3,
dataset='TopDownCocoDataset',
dataset_info=None,
show=False,
out_file=None):
"""Visualize the pose tracking results on the image.
Args:
model (nn.Module): The loaded detector.
img (str | np.ndarray): Image filename or loaded image.
result (list[dict]): The results to draw over `img`
(bbox_result, pose_result).
radius (int): Radius of circles.
thickness (int): Thickness of lines.
kpt_score_thr (float): The threshold to visualize the keypoints.
skeleton (list[tuple]): Default None.
show (bool): Whether to show the image. Default True.
out_file (str|None): The filename of the output visualization image.
"""
if hasattr(model, 'module'):
model = model.module
palette = np.array([[255, 128, 0], [255, 153, 51], [255, 178, 102],
[230, 230, 0], [255, 153, 255], [153, 204, 255],
[255, 102, 255], [255, 51, 255], [102, 178, 255],
[51, 153, 255], [255, 153, 153], [255, 102, 102],
[255, 51, 51], [153, 255, 153], [102, 255, 102],
[51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0],
[255, 255, 255]])
if dataset_info is None and dataset is not None:
warnings.warn(
'dataset is deprecated.'
'Please set `dataset_info` in the config.'
'Check https://github.com/open-mmlab/mmpose/pull/663 for details.',
DeprecationWarning)
# TODO: These will be removed in the later versions.
if dataset in ('TopDownCocoDataset', 'BottomUpCocoDataset',
'TopDownOCHumanDataset'):
kpt_num = 17
skeleton = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12],
[5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9],
[8, 10], [1, 2], [0, 1], [0, 2], [1, 3], [2, 4],
[3, 5], [4, 6]]
elif dataset == 'TopDownCocoWholeBodyDataset':
kpt_num = 133
skeleton = [[15, 13], [13, 11], [16, 14], [14, 12], [11, 12],
[5, 11], [6, 12], [5, 6], [5, 7], [6, 8], [7, 9],
[8, 10], [1, 2], [0, 1], [0, 2],
[1, 3], [2, 4], [3, 5], [4, 6], [15, 17], [15, 18],
[15, 19], [16, 20], [16, 21], [16, 22], [91, 92],
[92, 93], [93, 94], [94, 95], [91, 96], [96, 97],
[97, 98], [98, 99], [91, 100], [100, 101], [101, 102],
[102, 103], [91, 104], [104, 105], [105, 106],
[106, 107], [91, 108], [108, 109], [109, 110],
[110, 111], [112, 113], [113, 114], [114, 115],
[115, 116], [112, 117], [117, 118], [118, 119],
[119, 120], [112, 121], [121, 122], [122, 123],
[123, 124], [112, 125], [125, 126], [126, 127],
[127, 128], [112, 129], [129, 130], [130, 131],
[131, 132]]
radius = 1
elif dataset == 'TopDownAicDataset':
kpt_num = 14
skeleton = [[2, 1], [1, 0], [0, 13], [13, 3], [3, 4], [4, 5],
[8, 7], [7, 6], [6, 9], [9, 10], [10, 11], [12, 13],
[0, 6], [3, 9]]
elif dataset == 'TopDownMpiiDataset':
kpt_num = 16
skeleton = [[0, 1], [1, 2], [2, 6], [6, 3], [3, 4], [4, 5], [6, 7],
[7, 8], [8, 9], [8, 12], [12, 11], [11, 10], [8, 13],
[13, 14], [14, 15]]
elif dataset in ('OneHand10KDataset', 'FreiHandDataset',
'PanopticDataset'):
kpt_num = 21
skeleton = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7],
[7, 8], [0, 9], [9, 10], [10, 11], [11, 12], [0, 13],
[13, 14], [14, 15], [15, 16], [0, 17], [17, 18],
[18, 19], [19, 20]]
elif dataset == 'InterHand2DDataset':
kpt_num = 21
skeleton = [[0, 1], [1, 2], [2, 3], [4, 5], [5, 6], [6, 7], [8, 9],
[9, 10], [10, 11], [12, 13], [13, 14], [14, 15],
[16, 17], [17, 18], [18, 19], [3, 20], [7, 20],
[11, 20], [15, 20], [19, 20]]
else:
raise NotImplementedError()
elif dataset_info is not None:
kpt_num = dataset_info.keypoint_num
skeleton = dataset_info.skeleton
for res in result:
track_id = res['track_id']
bbox_color = palette[track_id % len(palette)]
pose_kpt_color = palette[[track_id % len(palette)] * kpt_num]
pose_link_color = palette[[track_id % len(palette)] * len(skeleton)]
img = model.show_result(
img, [res],
skeleton,
radius=radius,
thickness=thickness,
pose_kpt_color=pose_kpt_color,
pose_link_color=pose_link_color,
bbox_color=tuple(bbox_color.tolist()),
kpt_score_thr=kpt_score_thr,
show=show,
out_file=out_file)
return img
third_party/hamer/third-party/ViTPose/mmpose/apis/test.py
Vendored
+191
View File
@@ -0,0 +1,191 @@
# Copyright (c) OpenMMLab. All rights reserved.
import os.path as osp
import pickle
import shutil
import tempfile
import mmcv
import torch
import torch.distributed as dist
from mmcv.runner import get_dist_info
def single_gpu_test(model, data_loader):
"""Test model with a single gpu.
This method tests model with a single gpu and displays test progress bar.
Args:
model (nn.Module): Model to be tested.
data_loader (nn.Dataloader): Pytorch data loader.
Returns:
list: The prediction results.
"""
model.eval()
results = []
dataset = data_loader.dataset
prog_bar = mmcv.ProgressBar(len(dataset))
for data in data_loader:
with torch.no_grad():
result = model(return_loss=False, **data)
results.append(result)
# use the first key as main key to calculate the batch size
batch_size = len(next(iter(data.values())))
for _ in range(batch_size):
prog_bar.update()
return results
def multi_gpu_test(model, data_loader, tmpdir=None, gpu_collect=False):
"""Test model with multiple gpus.
This method tests model with multiple gpus and collects the results
under two different modes: gpu and cpu modes. By setting 'gpu_collect=True'
it encodes results to gpu tensors and use gpu communication for results
collection. On cpu mode it saves the results on different gpus to 'tmpdir'
and collects them by the rank 0 worker.
Args:
model (nn.Module): Model to be tested.
data_loader (nn.Dataloader): Pytorch data loader.
tmpdir (str): Path of directory to save the temporary results from
different gpus under cpu mode.
gpu_collect (bool): Option to use either gpu or cpu to collect results.
Returns:
list: The prediction results.
"""
model.eval()
results = []
dataset = data_loader.dataset
rank, world_size = get_dist_info()
if rank == 0:
prog_bar = mmcv.ProgressBar(len(dataset))
for data in data_loader:
with torch.no_grad():
result = model(return_loss=False, **data)
results.append(result)
if rank == 0:
# use the first key as main key to calculate the batch size
batch_size = len(next(iter(data.values())))
for _ in range(batch_size * world_size):
prog_bar.update()
# collect results from all ranks
if gpu_collect:
results = collect_results_gpu(results, len(dataset))
else:
results = collect_results_cpu(results, len(dataset), tmpdir)
return results
def collect_results_cpu(result_part, size, tmpdir=None):
"""Collect results in cpu mode.
It saves the results on different gpus to 'tmpdir' and collects
them by the rank 0 worker.
Args:
result_part (list): Results to be collected
size (int): Result size.
tmpdir (str): Path of directory to save the temporary results from
different gpus under cpu mode. Default: None
Returns:
list: Ordered results.
"""
rank, world_size = get_dist_info()
# create a tmp dir if it is not specified
if tmpdir is None:
MAX_LEN = 512
# 32 is whitespace
dir_tensor = torch.full((MAX_LEN, ),
32,
dtype=torch.uint8,
device='cuda')
if rank == 0:
mmcv.mkdir_or_exist('.dist_test')
tmpdir = tempfile.mkdtemp(dir='.dist_test')
tmpdir = torch.tensor(
bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda')
dir_tensor[:len(tmpdir)] = tmpdir
dist.broadcast(dir_tensor, 0)
tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip()
else:
mmcv.mkdir_or_exist(tmpdir)
# synchronizes all processes to make sure tmpdir exist
dist.barrier()
# dump the part result to the dir
mmcv.dump(result_part, osp.join(tmpdir, f'part_{rank}.pkl'))
# synchronizes all processes for loading pickle file
dist.barrier()
# collect all parts
if rank != 0:
return None
# load results of all parts from tmp dir
part_list = []
for i in range(world_size):
part_file = osp.join(tmpdir, f'part_{i}.pkl')
part_list.append(mmcv.load(part_file))
# sort the results
ordered_results = []
for res in zip(*part_list):
ordered_results.extend(list(res))
# the dataloader may pad some samples
ordered_results = ordered_results[:size]
# remove tmp dir
shutil.rmtree(tmpdir)
return ordered_results
def collect_results_gpu(result_part, size):
"""Collect results in gpu mode.
It encodes results to gpu tensors and use gpu communication for results
collection.
Args:
result_part (list): Results to be collected
size (int): Result size.
Returns:
list: Ordered results.
"""
rank, world_size = get_dist_info()
# dump result part to tensor with pickle
part_tensor = torch.tensor(
bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda')
# gather all result part tensor shape
shape_tensor = torch.tensor(part_tensor.shape, device='cuda')
shape_list = [shape_tensor.clone() for _ in range(world_size)]
dist.all_gather(shape_list, shape_tensor)
# padding result part tensor to max length
shape_max = torch.tensor(shape_list).max()
part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda')
part_send[:shape_tensor[0]] = part_tensor
part_recv_list = [
part_tensor.new_zeros(shape_max) for _ in range(world_size)
]
# gather all result part
dist.all_gather(part_recv_list, part_send)
if rank == 0:
part_list = []
for recv, shape in zip(part_recv_list, shape_list):
part_list.append(
pickle.loads(recv[:shape[0]].cpu().numpy().tobytes()))
# sort the results
ordered_results = []
for res in zip(*part_list):
ordered_results.extend(list(res))
# the dataloader may pad some samples
ordered_results = ordered_results[:size]
return ordered_results
return None
third_party/hamer/third-party/ViTPose/mmpose/apis/train.py
Vendored
+200
View File
@@ -0,0 +1,200 @@
# Copyright (c) OpenMMLab. All rights reserved.
import warnings
import mmcv
import numpy as np
import torch
import torch.distributed as dist
from mmcv.parallel import MMDataParallel, MMDistributedDataParallel
from mmcv.runner import (DistSamplerSeedHook, EpochBasedRunner, OptimizerHook,
get_dist_info)
from mmcv.utils import digit_version
from mmpose.core import DistEvalHook, EvalHook, build_optimizers
from mmpose.core.distributed_wrapper import DistributedDataParallelWrapper
from mmpose.datasets import build_dataloader, build_dataset
from mmpose.utils import get_root_logger
try:
from mmcv.runner import Fp16OptimizerHook
except ImportError:
warnings.warn(
'Fp16OptimizerHook from mmpose will be deprecated from '
'v0.15.0. Please install mmcv>=1.1.4', DeprecationWarning)
from mmpose.core import Fp16OptimizerHook
def init_random_seed(seed=None, device='cuda'):
"""Initialize random seed.
If the seed is not set, the seed will be automatically randomized,
and then broadcast to all processes to prevent some potential bugs.
Args:
seed (int, Optional): The seed. Default to None.
device (str): The device where the seed will be put on.
Default to 'cuda'.
Returns:
int: Seed to be used.
"""
if seed is not None:
return seed
# Make sure all ranks share the same random seed to prevent
# some potential bugs. Please refer to
# https://github.com/open-mmlab/mmdetection/issues/6339
rank, world_size = get_dist_info()
seed = np.random.randint(2**31)
if world_size == 1:
return seed
if rank == 0:
random_num = torch.tensor(seed, dtype=torch.int32, device=device)
else:
random_num = torch.tensor(0, dtype=torch.int32, device=device)
dist.broadcast(random_num, src=0)
return random_num.item()
def train_model(model,
dataset,
cfg,
distributed=False,
validate=False,
timestamp=None,
meta=None):
"""Train model entry function.
Args:
model (nn.Module): The model to be trained.
dataset (Dataset): Train dataset.
cfg (dict): The config dict for training.
distributed (bool): Whether to use distributed training.
Default: False.
validate (bool): Whether to do evaluation. Default: False.
timestamp (str | None): Local time for runner. Default: None.
meta (dict | None): Meta dict to record some important information.
Default: None
"""
logger = get_root_logger(cfg.log_level)
# prepare data loaders
dataset = dataset if isinstance(dataset, (list, tuple)) else [dataset]
# step 1: give default values and override (if exist) from cfg.data
loader_cfg = {
**dict(
seed=cfg.get('seed'),
drop_last=False,
dist=distributed,
num_gpus=len(cfg.gpu_ids)),
**({} if torch.__version__ != 'parrots' else dict(
prefetch_num=2,
pin_memory=False,
)),
**dict((k, cfg.data[k]) for k in [
'samples_per_gpu',
'workers_per_gpu',
'shuffle',
'seed',
'drop_last',
'prefetch_num',
'pin_memory',
'persistent_workers',
] if k in cfg.data)
}
# step 2: cfg.data.train_dataloader has highest priority
train_loader_cfg = dict(loader_cfg, **cfg.data.get('train_dataloader', {}))
data_loaders = [build_dataloader(ds, **train_loader_cfg) for ds in dataset]
# determine whether use adversarial training precess or not
use_adverserial_train = cfg.get('use_adversarial_train', False)
# put model on gpus
if distributed:
find_unused_parameters = cfg.get('find_unused_parameters', False)
# Sets the `find_unused_parameters` parameter in
# torch.nn.parallel.DistributedDataParallel
if use_adverserial_train:
# Use DistributedDataParallelWrapper for adversarial training
model = DistributedDataParallelWrapper(
model,
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
else:
model = MMDistributedDataParallel(
model.cuda(),
device_ids=[torch.cuda.current_device()],
broadcast_buffers=False,
find_unused_parameters=find_unused_parameters)
else:
if digit_version(mmcv.__version__) >= digit_version(
'1.4.4') or torch.cuda.is_available():
model = MMDataParallel(model, device_ids=cfg.gpu_ids)
else:
warnings.warn(
'We recommend to use MMCV >= 1.4.4 for CPU training. '
'See https://github.com/open-mmlab/mmpose/pull/1157 for '
'details.')
# build runner
optimizer = build_optimizers(model, cfg.optimizer)
runner = EpochBasedRunner(
model,
optimizer=optimizer,
work_dir=cfg.work_dir,
logger=logger,
meta=meta)
# an ugly workaround to make .log and .log.json filenames the same
runner.timestamp = timestamp
if use_adverserial_train:
# The optimizer step process is included in the train_step function
# of the model, so the runner should NOT include optimizer hook.
optimizer_config = None
else:
# fp16 setting
fp16_cfg = cfg.get('fp16', None)
if fp16_cfg is not None:
optimizer_config = Fp16OptimizerHook(
**cfg.optimizer_config, **fp16_cfg, distributed=distributed)
elif distributed and 'type' not in cfg.optimizer_config:
optimizer_config = OptimizerHook(**cfg.optimizer_config)
else:
optimizer_config = cfg.optimizer_config
# register hooks
runner.register_training_hooks(cfg.lr_config, optimizer_config,
cfg.checkpoint_config, cfg.log_config,
cfg.get('momentum_config', None))
if distributed:
runner.register_hook(DistSamplerSeedHook())
# register eval hooks
if validate:
eval_cfg = cfg.get('evaluation', {})
val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
dataloader_setting = dict(
samples_per_gpu=1,
workers_per_gpu=cfg.data.get('workers_per_gpu', 1),
# cfg.gpus will be ignored if distributed
num_gpus=len(cfg.gpu_ids),
dist=distributed,
drop_last=False,
shuffle=False)
dataloader_setting = dict(dataloader_setting,
**cfg.data.get('val_dataloader', {}))
val_dataloader = build_dataloader(val_dataset, **dataloader_setting)
eval_hook = DistEvalHook if distributed else EvalHook
runner.register_hook(eval_hook(val_dataloader, **eval_cfg))
if cfg.resume_from:
runner.resume(cfg.resume_from)
elif cfg.load_from:
runner.load_checkpoint(cfg.load_from)
runner.run(data_loaders, cfg.workflow, cfg.total_epochs)