Spaces:
Running
on
L40S
Running
on
L40S
import warnings | |
from typing import Iterable, List, Optional, Tuple, Union | |
import numpy as np | |
import detrsmpl.core.conventions.keypoints_mapping as keypoints_mapping | |
from detrsmpl.core.renderer.matplotlib3d_renderer import Axes3dJointsRenderer | |
from detrsmpl.utils.demo_utils import get_different_colors | |
from detrsmpl.utils.keypoint_utils import search_limbs | |
from detrsmpl.utils.path_utils import prepare_output_path | |
def _norm_pose(pose_numpy: np.ndarray, min_value: Union[float, int], | |
max_value: Union[float, int], mask: Union[np.ndarray, list]): | |
"""Normalize the poses and make the center close to axis center.""" | |
assert max_value > min_value | |
pose_np_normed = pose_numpy.copy() | |
if not mask: | |
mask = list(range(pose_numpy.shape[-2])) | |
axis_num = 3 | |
axis_stat = np.zeros(shape=[axis_num, 4]) | |
for axis_index in range(axis_num): | |
axis_data = pose_np_normed[..., mask, axis_index] | |
axis_min = np.min(axis_data) | |
axis_max = np.max(axis_data) | |
axis_mid = (axis_min + axis_max) / 2.0 | |
axis_span = axis_max - axis_min | |
axis_stat[axis_index] = np.asarray( | |
(axis_min, axis_max, axis_mid, axis_span)) | |
target_mid = (max_value + min_value) / 2.0 | |
max_span = np.max(axis_stat[:, 3]) | |
target_span = max_value - min_value | |
for axis_index in range(axis_num): | |
pose_np_normed[..., axis_index] = \ | |
pose_np_normed[..., axis_index] - \ | |
axis_stat[axis_index, 2] | |
pose_np_normed = pose_np_normed / max_span * target_span | |
pose_np_normed = pose_np_normed + target_mid | |
return pose_np_normed | |
def visualize_kp3d( | |
kp3d: np.ndarray, | |
output_path: Optional[str] = None, | |
limbs: Optional[Union[np.ndarray, List[int]]] = None, | |
palette: Optional[Iterable[int]] = None, | |
data_source: str = 'coco', | |
mask: Optional[Union[list, tuple, np.ndarray]] = None, | |
start: int = 0, | |
end: Optional[int] = None, | |
resolution: Union[list, Tuple[int, int]] = (1024, 1024), | |
fps: Union[float, int] = 30, | |
frame_names: Optional[Union[List[str], str]] = None, | |
orbit_speed: Union[float, int] = 0.5, | |
value_range: Union[Tuple[int, int], list] = (-100, 100), | |
pop_parts: Iterable[str] = (), | |
disable_limbs: bool = False, | |
return_array: Optional[bool] = None, | |
convention: str = 'opencv', | |
keypoints_factory: dict = keypoints_mapping.KEYPOINTS_FACTORY, | |
) -> Union[None, np.ndarray]: | |
"""Visualize 3d keypoints to a video with matplotlib. Support multi person | |
and specified limb connections. | |
Args: | |
kp3d (np.ndarray): shape could be (f * J * 4/3/2) or | |
(f * num_person * J * 4/3/2) | |
output_path (str): output video path image folder. | |
limbs (Optional[Union[np.ndarray, List[int]]], optional): | |
if not specified, the limbs will be searched by search_limbs, | |
this option is for free skeletons like BVH file. | |
Defaults to None. | |
palette (Iterable, optional): specified palette, three int represents | |
(B, G, R). Should be tuple or list. | |
Defaults to None. | |
data_source (str, optional): data source type. Defaults to 'coco'. | |
choose in ['coco', 'smplx', 'smpl', 'coco_wholebody', | |
'mpi_inf_3dhp', 'mpi_inf_3dhp_test', 'h36m', 'pw3d', 'mpii'] | |
mask (Optional[Union[list, tuple, np.ndarray]], optional): | |
mask to mask out the incorrect points. Defaults to None. | |
start (int, optional): start frame index. Defaults to 0. | |
end (int, optional): end frame index. | |
Could be positive int or negative int or None. | |
None represents include all the frames. | |
Defaults to None. | |
resolution (Union[list, Tuple[int, int]], optional): | |
(width, height) of the output video | |
will be the same size as the original images if not specified. | |
Defaults to None. | |
fps (Union[float, int], optional): fps. Defaults to 30. | |
frame_names (Optional[Union[List[str], str]], optional): List(should be | |
the same as frame numbers) or single string or string format | |
(like 'frame%06d')for frame title, no title if None. | |
Defaults to None. | |
orbit_speed (Union[float, int], optional): orbit speed of camera. | |
Defaults to 0.5. | |
value_range (Union[Tuple[int, int], list], optional): | |
range of axis value. Defaults to (-100, 100). | |
pop_parts (Iterable[str], optional): The body part names you do not | |
want to visualize. Choose in ['left_eye','right_eye', 'nose', | |
'mouth', 'face', 'left_hand', 'right_hand']Defaults to []. | |
disable_limbs (bool, optional): whether need to disable drawing limbs. | |
Defaults to False. | |
return_array (bool, optional): Whether to return images as opencv array | |
.If None, an array will be returned when frame number is below 100. | |
Defaults to None. | |
keypoints_factory (dict, optional): Dict of all the conventions. | |
Defaults to KEYPOINTS_FACTORY. | |
Raises: | |
TypeError: check the type of input keypoints. | |
FileNotFoundError: check the output video path. | |
Returns: | |
Union[None, np.ndarray]. | |
""" | |
# check input shape | |
if not isinstance(kp3d, np.ndarray): | |
raise TypeError( | |
f'Input type is {type(kp3d)}, which should be numpy.ndarray.') | |
kp3d = kp3d.copy() | |
if kp3d.shape[-1] == 2: | |
kp3d = np.concatenate([kp3d, np.zeros_like(kp3d)[..., 0:1]], axis=-1) | |
warnings.warn( | |
'The input array is 2-Dimensional coordinates, will concatenate ' + | |
f'zeros to the last axis. The new array shape: {kp3d.shape}') | |
elif kp3d.shape[-1] >= 4: | |
kp3d = kp3d[..., :3] | |
warnings.warn( | |
'The input array has more than 3-Dimensional coordinates, will ' + | |
'keep only the first 3-Dimensions of the last axis. The new ' + | |
f'array shape: {kp3d.shape}') | |
if kp3d.ndim == 3: | |
kp3d = np.expand_dims(kp3d, 1) | |
num_frames = kp3d.shape[0] | |
assert kp3d.ndim == 4 | |
assert kp3d.shape[-1] == 3 | |
if return_array is None: | |
if num_frames > 100: | |
return_array = False | |
else: | |
return_array = True | |
# check data_source & mask | |
if data_source not in keypoints_factory: | |
raise ValueError('Wrong data_source. Should choose in' + | |
f'{list(keypoints_factory.keys())}') | |
if mask is not None: | |
if not isinstance(mask, np.ndarray): | |
mask = np.array(mask).reshape(-1) | |
assert mask.shape == ( | |
len(keypoints_factory[data_source]), | |
), f'mask length should fit with keypoints number \ | |
{len(keypoints_factory[data_source])}' | |
# check the output path | |
if output_path is not None: | |
prepare_output_path(output_path, | |
path_type='auto', | |
tag='output video', | |
allowed_suffix=['.mp4', '.gif', '']) | |
# slice the frames | |
end = num_frames if end is None else end | |
kp3d = kp3d[start:end] | |
# norm the coordinates | |
if value_range is not None: | |
# norm pose location to value_range (70% value range) | |
mask_index = np.where(np.array(mask) > 0) if mask is not None else None | |
margin_width = abs(value_range[1] - value_range[1]) * 0.15 | |
pose_np_normed = _norm_pose(kp3d, value_range[0] + margin_width, | |
value_range[1] - margin_width, mask_index) | |
input_pose_np = pose_np_normed | |
else: | |
input_pose_np = kp3d | |
# determine the limb connections and palettes | |
if limbs is not None: | |
limbs_target, limbs_palette = { | |
'body': limbs.tolist() if isinstance(limbs, np.ndarray) else limbs | |
}, get_different_colors(len(limbs)) | |
else: | |
limbs_target, limbs_palette = search_limbs(data_source=data_source, | |
mask=mask) | |
if palette is not None: | |
limbs_palette = np.array(palette, dtype=np.uint8)[None] | |
# check and pop the pop_parts | |
assert set(pop_parts).issubset( | |
keypoints_mapping.human_data.HUMAN_DATA_PALETTE.keys( | |
)), f'wrong part_names in pop_parts, could only \ | |
choose in{set(keypoints_mapping.human_data.HUMAN_DATA_PALETTE.keys())}' | |
for part_name in pop_parts: | |
if part_name in limbs_target: | |
limbs_target.pop(part_name) | |
# initialize renderer and start render | |
renderer = Axes3dJointsRenderer() | |
renderer.init_camera(cam_hori_speed=orbit_speed, cam_elev_speed=0.2) | |
renderer.set_connections(limbs_target, limbs_palette) | |
if isinstance(frame_names, str): | |
if '%' in frame_names: | |
frame_names = [ | |
frame_names % index for index in range(input_pose_np.shape[0]) | |
] | |
else: | |
frame_names = [frame_names] * input_pose_np.shape[0] | |
image_array = renderer.render_kp3d_to_video(input_pose_np, | |
output_path, | |
convention, | |
fps=fps, | |
resolution=resolution, | |
visual_range=value_range, | |
frame_names=frame_names, | |
disable_limbs=disable_limbs, | |
return_array=return_array) | |
return image_array | |