import colorsys
import os
from collections import defaultdict
from contextlib import contextmanager
from functools import partial
from pathlib import Path
import mmcv
import numpy as np
from mmcv import Timer
from scipy import interpolate
from detrsmpl.core.post_processing import build_post_processing
try:
from typing import Literal
except ImportError:
from typing_extensions import Literal
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])
"""
if not isinstance(bbox_xyxy, np.ndarray):
raise TypeError(
f'Input type is {type(bbox_xyxy)}, which should be numpy.ndarray.')
bbox_xywh = bbox_xyxy.copy()
bbox_xywh[..., 2] = bbox_xywh[..., 2] - bbox_xywh[..., 0]
bbox_xywh[..., 3] = bbox_xywh[..., 3] - bbox_xywh[..., 1]
return bbox_xywh
def xywh2xyxy(bbox_xywh):
"""Transform the bbox format from xywh to x1y1x2y2.
Args:
bbox_xywh (np.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])
"""
if not isinstance(bbox_xywh, np.ndarray):
raise TypeError(
f'Input type is {type(bbox_xywh)}, which should be numpy.ndarray.')
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(bbox_xywh, aspect_ratio=1.0, bbox_scale_factor=1.25):
"""Convert xywh coordinates to center and scale.
Args:
bbox_xywh (numpy.ndarray): the height of the bbox_xywh
aspect_ratio (int, optional): Defaults to 1.0
bbox_scale_factor (float, optional): Defaults to 1.25
Returns:
numpy.ndarray: center of the bbox
numpy.ndarray: the scale of the bbox w & h
"""
if not isinstance(bbox_xywh, np.ndarray):
raise TypeError(
f'Input type is {type(bbox_xywh)}, which should be numpy.ndarray.')
bbox_xywh = bbox_xywh.copy()
pixel_std = 1
center = np.stack([
bbox_xywh[..., 0] + bbox_xywh[..., 2] * 0.5,
bbox_xywh[..., 1] + bbox_xywh[..., 3] * 0.5
], -1)
mask_h = bbox_xywh[..., 2] > aspect_ratio * bbox_xywh[..., 3]
mask_w = ~mask_h
bbox_xywh[mask_h, 3] = bbox_xywh[mask_h, 2] / aspect_ratio
bbox_xywh[mask_w, 2] = bbox_xywh[mask_w, 3] * aspect_ratio
scale = np.stack([
bbox_xywh[..., 2] * 1.0 / pixel_std,
bbox_xywh[..., 3] * 1.0 / pixel_std
], -1)
scale = scale * bbox_scale_factor
return center, scale
def convert_crop_cam_to_orig_img(cam: np.ndarray,
bbox: np.ndarray,
img_width: int,
img_height: int,
aspect_ratio: float = 1.0,
bbox_scale_factor: float = 1.25,
bbox_format: Literal['xyxy', 'xywh',
'cs'] = 'xyxy'):
"""This function is modified from [VIBE](https://github.com/
mkocabas/VIBE/blob/master/lib/utils/demo_utils.py#L242-L259). Original
license please see docs/additional_licenses.md.
Args:
cam (np.ndarray): cam (ndarray, shape=(frame, 3) or
(frame,num_person, 3)):
weak perspective camera in cropped img coordinates
bbox (np.ndarray): bbox coordinates
img_width (int): original image width
img_height (int): original image height
aspect_ratio (float, optional): Defaults to 1.0.
bbox_scale_factor (float, optional): Defaults to 1.25.
bbox_format (Literal['xyxy', 'xywh', 'cs']): Defaults to 'xyxy'.
'xyxy' means the left-up point and right-bottomn point of the
bbox.
'xywh' means the left-up point and the width and height of the
bbox.
'cs' means the center of the bbox (x,y) and the scale of the
bbox w & h.
Returns:
orig_cam: shape = (frame, 4) or (frame, num_person, 4)
"""
if not isinstance(bbox, np.ndarray):
raise TypeError(
f'Input type is {type(bbox)}, which should be numpy.ndarray.')
bbox = bbox.copy()
if bbox_format == 'xyxy':
bbox_xywh = xyxy2xywh(bbox)
center, scale = box2cs(bbox_xywh, aspect_ratio, bbox_scale_factor)
bbox_cs = np.concatenate([center, scale], axis=-1)
elif bbox_format == 'xywh':
center, scale = box2cs(bbox, aspect_ratio, bbox_scale_factor)
bbox_cs = np.concatenate([center, scale], axis=-1)
elif bbox_format == 'cs':
bbox_cs = bbox
else:
raise ValueError('Only supports the format of `xyxy`, `cs` and `xywh`')
cx, cy, h = bbox_cs[..., 0], bbox_cs[..., 1], bbox_cs[..., 2] + 1e-6
hw, hh = img_width / 2., img_height / 2.
sx = cam[..., 0] * (1. / (img_width / h))
sy = cam[..., 0] * (1. / (img_height / h))
tx = ((cx - hw) / hw / (sx + 1e-6)) + cam[..., 1]
ty = ((cy - hh) / hh / (sy + 1e-6)) + cam[..., 2]
orig_cam = np.stack([sx, sy, tx, ty], axis=-1)
return orig_cam
def convert_bbox_to_intrinsic(bboxes: np.ndarray,
img_width: int = 224,
img_height: int = 224,
bbox_scale_factor: float = 1.25,
bbox_format: Literal['xyxy', 'xywh'] = 'xyxy'):
"""Convert bbox to intrinsic parameters.
Args:
bbox (np.ndarray): (frame, num_person, 4), (frame, 4), or (4,)
img_width (int): image width of training data.
img_height (int): image height of training data.
bbox_scale_factor (float): scale factor for expanding the bbox.
bbox_format (Literal['xyxy', 'xywh'] ): 'xyxy' means the left-up point
and right-bottomn point of the bbox.
'xywh' means the left-up point and the width and height of the
bbox.
Returns:
np.ndarray: (frame, num_person, 3, 3), (frame, 3, 3) or (3,3)
"""
if not isinstance(bboxes, np.ndarray):
raise TypeError(
f'Input type is {type(bboxes)}, which should be numpy.ndarray.')
assert bbox_format in ['xyxy', 'xywh']
if bbox_format == 'xyxy':
bboxes = xyxy2xywh(bboxes)
center_x = bboxes[..., 0] + bboxes[..., 2] / 2.0
center_y = bboxes[..., 1] + bboxes[..., 3] / 2.0
W = np.max(bboxes[..., 2:], axis=-1) * bbox_scale_factor
num_frame = bboxes.shape[0]
if bboxes.ndim == 3:
num_person = bboxes.shape[1]
Ks = np.zeros((num_frame, num_person, 3, 3))
elif bboxes.ndim == 2:
Ks = np.zeros((num_frame, 3, 3))
elif bboxes.ndim == 1:
Ks = np.zeros((3, 3))
else:
raise ValueError('Wrong input bboxes shape {bboxes.shape}')
Ks[..., 0, 0] = W / img_width
Ks[..., 1, 1] = W / img_height
Ks[..., 0, 2] = center_x - W / 2.0
Ks[..., 1, 2] = center_y - W / 2.0
Ks[..., 2, 2] = 1
return Ks
def get_default_hmr_intrinsic(num_frame=1,
focal_length=1000,
det_width=224,
det_height=224) -> np.ndarray:
"""Get default hmr intrinsic, defined by how you trained.
Args:
num_frame (int, optional): num of frames. Defaults to 1.
focal_length (int, optional): defined same as your training.
Defaults to 1000.
det_width (int, optional): the size you used to detect.
Defaults to 224.
det_height (int, optional): the size you used to detect.
Defaults to 224.
Returns:
np.ndarray: shape of (N, 3, 3)
"""
K = np.zeros((num_frame, 3, 3))
K[:, 0, 0] = focal_length
K[:, 1, 1] = focal_length
K[:, 0, 2] = det_width / 2
K[:, 1, 2] = det_height / 2
K[:, 2, 2] = 1
return K
def convert_kp2d_to_bbox(
kp2d: np.ndarray,
bbox_format: Literal['xyxy', 'xywh'] = 'xyxy') -> np.ndarray:
"""Convert kp2d to bbox.
Args:
kp2d (np.ndarray): shape should be (num_frame, num_points, 2/3)
or (num_frame, num_person, num_points, 2/3).
bbox_format (Literal['xyxy', 'xywh'], optional): Defaults to 'xyxy'.
Returns:
np.ndarray: shape will be (num_frame, num_person, 4)
"""
assert bbox_format in ['xyxy', 'xywh']
if kp2d.ndim == 2:
kp2d = kp2d[None, None]
elif kp2d.ndim == 3:
kp2d = kp2d[:, None]
num_frame, num_person, _, _ = kp2d.shape
x1 = np.max(kp2d[..., 0], axis=-2)
y1 = np.max(kp2d[..., 1], axis=-2)
x2 = np.max(kp2d[..., 2], axis=-2)
y2 = np.max(kp2d[..., 3], axis=-2)
bbox = np.concatenate([x1, y1, x2, y2], axis=-1)
assert bbox.shape == (num_frame, num_person, 4)
if bbox_format == 'xywh':
bbox = xyxy2xywh(bbox)
return bbox
def convert_verts_to_cam_coord(verts,
pred_cams,
bboxes_xy,
focal_length=5000.,
bbox_scale_factor=1.25,
bbox_format='xyxy'):
"""Convert vertices from the world coordinate to camera coordinate.
Args:
verts ([np.ndarray]): The vertices in the world coordinate.
The shape is (frame,num_person,6890,3), (frame,6890,3),
or (6890,3).
pred_cams ([np.ndarray]): Camera parameters estimated by HMR or SPIN.
The shape is (frame,num_person,3), (frame,3), or (3,).
bboxes_xy ([np.ndarray]): (frame, num_person, 4|5), (frame, 4|5),
or (4|5,)
focal_length ([float],optional): Defined same as your training.
bbox_scale_factor (float): scale factor for expanding the bbox.
bbox_format (Literal['xyxy', 'xywh'] ): 'xyxy' means the left-up point
and right-bottomn point of the bbox.
'xywh' means the left-up point and the width and height of the
bbox.
Returns:
np.ndarray: The vertices in the camera coordinate.
The shape is (frame,num_person,6890,3) or (frame,6890,3).
np.ndarray: The intrinsic parameters of the pred_cam.
The shape is (num_frame, 3, 3).
"""
K0 = get_default_hmr_intrinsic(focal_length=focal_length,
det_height=224,
det_width=224)
K1 = convert_bbox_to_intrinsic(bboxes_xy,
bbox_scale_factor=bbox_scale_factor,
bbox_format=bbox_format)
# K1K0(RX+T)-> K0(K0_inv K1K0)
Ks = np.linalg.inv(K0) @ K1 @ K0
# convert vertices from world to camera
cam_trans = np.concatenate([
pred_cams[..., [1]], pred_cams[..., [2]], 2 * focal_length /
(224 * pred_cams[..., [0]] + 1e-9)
], -1)
verts = verts + cam_trans[..., None, :]
if verts.ndim == 4:
verts = np.einsum('fnij,fnkj->fnki', Ks, verts)
elif verts.ndim == 3:
verts = np.einsum('fij,fkj->fki', Ks, verts)
elif verts.ndim == 2:
verts = np.einsum('fij,fkj->fki', Ks, verts[None])
return verts, K0
def smooth_process(x,
smooth_type='savgol',
cfg_base_dir='configs/_base_/post_processing/'):
"""Smooth the array with the specified smoothing type.
Args:
x (np.ndarray): Shape should be (frame,num_person,K,C)
or (frame,K,C).
smooth_type (str, optional): Smooth type.
choose in ['oneeuro', 'gaus1d', 'savgol','smoothnet',
'smoothnet_windowsize8','smoothnet_windowsize16',
'smoothnet_windowsize32','smoothnet_windowsize64'].
Defaults to 'savgol'. 'smoothnet' is default with windowsize=8.
cfg_base_dir (str, optional): Config base dir,
default configs/_base_/post_processing/
Raises:
ValueError: check the input smoothing type.
Returns:
np.ndarray: Smoothed data. The shape should be
(frame,num_person,K,C) or (frame,K,C).
"""
if smooth_type == 'smoothnet':
smooth_type = 'smoothnet_windowsize8'
assert smooth_type in [
'oneeuro', 'gaus1d', 'savgol', 'smoothnet_windowsize8',
'smoothnet_windowsize16', 'smoothnet_windowsize32',
'smoothnet_windowsize64'
]
cfg = os.path.join(cfg_base_dir, smooth_type + '.py')
if isinstance(cfg, str):
cfg = mmcv.Config.fromfile(cfg)
elif not isinstance(cfg, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
f'but got {type(cfg)}')
x = x.copy()
assert x.ndim == 3 or x.ndim == 4
smooth_func = build_post_processing(dict(cfg['smooth_cfg']))
if x.ndim == 4:
for i in range(x.shape[1]):
x[:, i] = smooth_func(x[:, i])
elif x.ndim == 3:
x = smooth_func(x)
return x
def speed_up_process(x,
speed_up_type='deciwatch',
cfg_base_dir='configs/_base_/post_processing/'):
"""Speed up the process with the specified speed up type.
Args:
x (np.ndarray): Shape should be (frame,num_person,K,C)
or (frame,K,C).
speed_up_type (str, optional): Speed up type.
choose in ['deciwatch',
'deciwatch_interval5_q1',
'deciwatch_interval5_q2',
'deciwatch_interval5_q3',
'deciwatch_interval5_q4',
'deciwatch_interval5_q5',
'deciwatch_interval10_q1',
'deciwatch_interval10_q2',
'deciwatch_interval10_q3',
'deciwatch_interval10_q4',
'deciwatch_interval10_q5',]. Defaults to 'deciwatch'.
cfg_base_dir (str, optional): Config base dir.
Defaults to 'configs/_base_/post_processing/'
Raises:
ValueError: check the input speed up type.
Returns:
np.ndarray: Completed data. The shape should be
(frame,num_person,K,C) or (frame,K,C).
"""
if speed_up_type == 'deciwatch':
speed_up_type = 'deciwatch_interval5_q3'
assert speed_up_type in [
'deciwatch_interval5_q1',
'deciwatch_interval5_q2',
'deciwatch_interval5_q3',
'deciwatch_interval5_q4',
'deciwatch_interval5_q5',
'deciwatch_interval10_q1',
'deciwatch_interval10_q2',
'deciwatch_interval10_q3',
'deciwatch_interval10_q4',
'deciwatch_interval10_q5',
]
cfg = os.path.join(cfg_base_dir, speed_up_type + '.py')
if isinstance(cfg, str):
cfg = mmcv.Config.fromfile(cfg)
elif not isinstance(cfg, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
f'but got {type(cfg)}')
x = x.clone()
assert x.ndim == 4 or x.ndim == 5
cfg_dict = cfg['speed_up_cfg']
cfg_dict['device'] = x.device
speed_up_func = build_post_processing(cfg_dict)
if x.ndim == 5:
for i in range(x.shape[1]):
x[:, i] = speed_up_func(x[:, i])
elif x.ndim == 4:
x = speed_up_func(x)
return np.array(x.cpu())
def get_speed_up_interval(speed_up_type,
cfg_base_dir='configs/_base_/post_processing/'):
"""Get the interval of specific speed up type.
Args:
speed_up_type (str, optional): Speed up type.
choose in ['deciwatch',
'deciwatch_interval5_q1',
'deciwatch_interval5_q2',
'deciwatch_interval5_q3',
'deciwatch_interval5_q4',
'deciwatch_interval5_q5',
'deciwatch_interval10_q1',
'deciwatch_interval10_q2',
'deciwatch_interval10_q3',
'deciwatch_interval10_q4',
'deciwatch_interval10_q5',]. Defaults to 'deciwatch'.
cfg_base_dir (str, optional): Config base dir,
default configs/_base_/post_processing/
Raises:
ValueError: check the input speed up type.
Returns:
int: speed up interval
"""
if speed_up_type == 'deciwatch':
speed_up_type = 'deciwatch_interval5_q3'
assert speed_up_type in [
'deciwatch_interval5_q1',
'deciwatch_interval5_q2',
'deciwatch_interval5_q3',
'deciwatch_interval5_q4',
'deciwatch_interval5_q5',
'deciwatch_interval10_q1',
'deciwatch_interval10_q2',
'deciwatch_interval10_q3',
'deciwatch_interval10_q4',
'deciwatch_interval10_q5',
]
cfg = os.path.join(cfg_base_dir, speed_up_type + '.py')
if isinstance(cfg, str):
cfg = mmcv.Config.fromfile(cfg)
elif not isinstance(cfg, mmcv.Config):
raise TypeError('config must be a filename or Config object, '
f'but got {type(cfg)}')
return cfg['speed_up_cfg']['interval']
def speed_up_interpolate(selected_frames, speed_up_frames, smpl_poses,
smpl_betas, pred_cams, bboxes_xyxy):
"""Interpolate smpl_betas, pred_cams, and bboxes_xyxyx for speed up.
Args:
selected_frames (np.ndarray): Shape should be (selected frame number).
speed_up_frames (int): Total speed up frame number
smpl_poses (np.ndarray): selected frame smpl poses parameter
smpl_betas (np.ndarray): selected frame smpl shape paeameter
pred_cams (np.ndarray): selected frame camera parameter
bboxes_xyxy (np.ndarray): selected frame bbox
Returns:
smpl_poses (np.ndarray): interpolated frame smpl poses parameter
smpl_betas (np.ndarray): interpolated frame smpl shape paeameter
pred_cams (np.ndarray): interpolated frame camera parameter
bboxes_xyxy (np.ndarray): interpolated frame bbox
"""
selected_frames = selected_frames[selected_frames <= speed_up_frames]
pred_cams[:speed_up_frames, :] = interpolate.interp1d(
selected_frames, pred_cams[selected_frames, :], kind='linear',
axis=0)(np.arange(0, max(selected_frames)))
bboxes_xyxy[:speed_up_frames, :] = interpolate.interp1d(
selected_frames,
bboxes_xyxy[selected_frames, :],
kind='linear',
axis=0)(np.arange(0, max(selected_frames)))
smpl_betas[:speed_up_frames, :] = interpolate.interp1d(
selected_frames, smpl_betas[selected_frames, :], kind='linear',
axis=0)(np.arange(0, max(selected_frames)))
return smpl_poses, smpl_betas, pred_cams, bboxes_xyxy
def process_mmtracking_results(mmtracking_results,
max_track_id,
bbox_thr=None):
"""Process mmtracking results.
Args:
mmtracking_results ([list]): mmtracking_results.
bbox_thr (float): threshold for bounding boxes.
max_track_id (int): the maximum track id.
Returns:
person_results ([list]): a list of tracked bounding boxes
max_track_id (int): the maximum track id.
instance_num (int): the number of instance.
"""
person_results = []
# 'track_results' is changed to 'track_bboxes'
# in https://github.com/open-mmlab/mmtracking/pull/300
if 'track_bboxes' in mmtracking_results:
tracking_results = mmtracking_results['track_bboxes'][0]
elif 'track_results' in mmtracking_results:
tracking_results = mmtracking_results['track_results'][0]
tracking_results = np.array(tracking_results)
if bbox_thr is not None:
assert tracking_results.shape[-1] == 6
valid_idx = np.where(tracking_results[:, 5] > bbox_thr)[0]
tracking_results = tracking_results[valid_idx]
for track in tracking_results:
person = {}
person['track_id'] = int(track[0])
if max_track_id < int(track[0]):
max_track_id = int(track[0])
person['bbox'] = track[1:]
person_results.append(person)
person_results = sorted(person_results, key=lambda x: x.get('track_id', 0))
instance_num = len(person_results)
return person_results, max_track_id, instance_num
def process_mmdet_results(mmdet_results, cat_id=1, bbox_thr=None):
"""Process mmdet results, and return a list of bboxes.
Args:
mmdet_results (list|tuple): mmdet results.
bbox_thr (float): threshold for bounding boxes.
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 = []
bboxes = np.array(bboxes)
if bbox_thr is not None:
assert bboxes.shape[-1] == 5
valid_idx = np.where(bboxes[:, 4] > bbox_thr)[0]
bboxes = bboxes[valid_idx]
for bbox in bboxes:
person = {}
person['bbox'] = bbox
person_results.append(person)
return person_results
def prepare_frames(input_path=None):
"""Prepare frames from input_path.
Args:
input_path (str, optional): Defaults to None.
Raises:
ValueError: check the input path.
Returns:
List[np.ndarray]: prepared frames
"""
if Path(input_path).is_file():
img_list = [mmcv.imread(input_path)]
if img_list[0] is None:
video = mmcv.VideoReader(input_path)
assert video.opened, f'Failed to load file {input_path}'
img_list = list(video)
elif Path(input_path).is_dir():
# input_type = 'folder'
file_list = [
os.path.join(input_path, fn) for fn in os.listdir(input_path)
if fn.lower().endswith(('.png', '.jpg'))
]
file_list.sort()
img_list = [mmcv.imread(img_path) for img_path in file_list]
assert len(img_list), f'Failed to load image from {input_path}'
else:
raise ValueError('Input path should be an file or folder.'
f' Got invalid input path: {input_path}')
return img_list
def extract_feature_sequence(extracted_results,
frame_idx,
causal,
seq_len,
step=1):
"""Extract the target frame from person results, and pad the sequence to a
fixed length.
Args:
extracted_results (List[List[Dict]]): Multi-frame feature extraction
results stored in a nested list. Each element of the outer list
is the feature extraction results of a single frame, and each
element of the inner list is the feature information of one person,
which contains:
features (ndarray): extracted features
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 first 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 feature extraction results stored in a
nested list with a length of seq_len.
int: The target frame index in the padded sequence.
"""
if causal:
frames_left = 0
frames_right = seq_len - 1
else:
frames_left = (seq_len - 1) // 2
frames_right = frames_left
num_frames = len(extracted_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)
extracted_results_seq = [extracted_results[0]] * pad_left + \
extracted_results[start:end:step] + [extracted_results[-1]] * pad_right
return extracted_results_seq
def get_different_colors(number_of_colors,
flag=0,
alpha: float = 1.0,
mode: str = 'bgr',
int_dtype: bool = True):
"""Get a numpy of colors of shape (N, 3)."""
mode = mode.lower()
assert set(mode).issubset({'r', 'g', 'b', 'a'})
nst0 = np.random.get_state()
np.random.seed(flag)
colors = []
for i in np.arange(0., 360., 360. / number_of_colors):
hue = i / 360.
lightness = (50 + np.random.rand() * 10) / 100.
saturation = (90 + np.random.rand() * 10) / 100.
colors.append(colorsys.hls_to_rgb(hue, lightness, saturation))
colors_np = np.asarray(colors)
if int_dtype:
colors_bgr = (255 * colors_np).astype(np.uint8)
else:
colors_bgr = colors_np.astype(np.float32)
# recover the random state
np.random.set_state(nst0)
color_dict = {}
if 'a' in mode:
color_dict['a'] = np.ones((colors_bgr.shape[0], 3)) * alpha
color_dict['b'] = colors_bgr[:, 0:1]
color_dict['g'] = colors_bgr[:, 1:2]
color_dict['r'] = colors_bgr[:, 2:3]
colors_final = []
for channel in mode:
colors_final.append(color_dict[channel])
colors_final = np.concatenate(colors_final, -1)
return colors_final
class RunningAverage():
r"""A helper class to calculate running average in a sliding window.
Args:
window (int): The size of the sliding window.
"""
def __init__(self, window: int = 1):
self.window = window
self._data = []
def update(self, value):
"""Update a new data sample."""
self._data.append(value)
self._data = self._data[-self.window:]
def average(self):
"""Get the average value of current window."""
return np.mean(self._data)
class StopWatch:
r"""A helper class to measure FPS and detailed time consuming of each phase
in a video processing loop or similar scenarios.
Args:
window (int): The sliding window size to calculate the running average
of the time consuming.
Example:
>>> from mmpose.utils import StopWatch
>>> import time
>>> stop_watch = StopWatch(window=10)
>>> with stop_watch.timeit('total'):
>>> time.sleep(0.1)
>>> # 'timeit' support nested use
>>> with stop_watch.timeit('phase1'):
>>> time.sleep(0.1)
>>> with stop_watch.timeit('phase2'):
>>> time.sleep(0.2)
>>> time.sleep(0.2)
>>> report = stop_watch.report()
"""
def __init__(self, window=1):
self.window = window
self._record = defaultdict(partial(RunningAverage, window=self.window))
self._timer_stack = []
@contextmanager
def timeit(self, timer_name='_FPS_'):
"""Timing a code snippet with an assigned name.
Args:
timer_name (str): The unique name of the interested code snippet to
handle multiple timers and generate reports. Note that '_FPS_'
is a special key that the measurement will be in `fps` instead
of `millisecond`. Also see `report` and `report_strings`.
Default: '_FPS_'.
Note:
This function should always be used in a `with` statement, as shown
in the example.
"""
self._timer_stack.append((timer_name, Timer()))
try:
yield
finally:
timer_name, timer = self._timer_stack.pop()
self._record[timer_name].update(timer.since_start())
def report(self, key=None):
"""Report timing information.
Returns:
dict: The key is the timer name and the value is the \
corresponding average time consuming.
"""
result = {
name: r.average() * 1000.
for name, r in self._record.items()
}
if '_FPS_' in result:
result['_FPS_'] = 1000. / result.pop('_FPS_')
if key is None:
return result
return result[key]
def report_strings(self):
"""Report timing information in texture strings.
Returns:
list(str): Each element is the information string of a timed \
event, in format of '{timer_name}: {time_in_ms}'. \
Specially, if timer_name is '_FPS_', the result will \
be converted to fps.
"""
result = self.report()
strings = []
if '_FPS_' in result:
strings.append(f'FPS: {result["_FPS_"]:>5.1f}')
strings += [f'{name}: {val:>3.0f}' for name, val in result.items()]
return strings
def reset(self):
self._record = defaultdict(list)
self._active_timer_stack = []