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

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-09-15 14:13:49 +08:00
commit 7e4ef6f98b
7002 changed files with 1452249 additions and 0 deletions
@@ -0,0 +1,4 @@
# Copyright (c) OpenMMLab. All rights reserved.
from .webcam_runner import WebcamRunner
__all__ = ['WebcamRunner']
@@ -0,0 +1,18 @@
# Copyright (c) OpenMMLab. All rights reserved.
from .builder import NODES
from .faceswap_node import FaceSwapNode
from .frame_effect_node import (BackgroundNode, BugEyeNode, MoustacheNode,
NoticeBoardNode, PoseVisualizerNode,
SaiyanNode, SunglassesNode)
from .helper_node import ModelResultBindingNode, MonitorNode, RecorderNode
from .mmdet_node import DetectorNode
from .mmpose_node import TopDownPoseEstimatorNode
from .valentinemagic_node import ValentineMagicNode
from .xdwendwen_node import XDwenDwenNode
__all__ = [
'NODES', 'PoseVisualizerNode', 'DetectorNode', 'TopDownPoseEstimatorNode',
'MonitorNode', 'BugEyeNode', 'SunglassesNode', 'ModelResultBindingNode',
'NoticeBoardNode', 'RecorderNode', 'FaceSwapNode', 'MoustacheNode',
'SaiyanNode', 'BackgroundNode', 'XDwenDwenNode', 'ValentineMagicNode'
]
@@ -0,0 +1,4 @@
# Copyright (c) OpenMMLab. All rights reserved.
from mmcv.utils import Registry
NODES = Registry('node')
@@ -0,0 +1,254 @@
# Copyright (c) OpenMMLab. All rights reserved.
from enum import IntEnum
from typing import List, Union
import cv2
import numpy as np
from mmpose.datasets import DatasetInfo
from .builder import NODES
from .frame_drawing_node import FrameDrawingNode
class Mode(IntEnum):
NONE = 0,
SHUFFLE = 1,
CLONE = 2
@NODES.register_module()
class FaceSwapNode(FrameDrawingNode):
def __init__(
self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
mode_key: Union[str, int],
):
super().__init__(name, frame_buffer, output_buffer, enable=True)
self.mode_key = mode_key
self.mode_index = 0
self.register_event(
self.mode_key, is_keyboard=True, handler_func=self.switch_mode)
self.history = dict(mode=None)
self._mode = Mode.SHUFFLE
@property
def mode(self):
return self._mode
def switch_mode(self):
"""Switch modes by updating mode index."""
self._mode = Mode((self._mode + 1) % len(Mode))
def draw(self, frame_msg):
if self.mode == Mode.NONE:
self.history = {'mode': Mode.NONE}
return frame_msg.get_image()
# Init history
if self.history['mode'] != self.mode:
self.history = {'mode': self.mode, 'target_map': {}}
# Merge pose results
pose_preds = self._merge_pose_results(frame_msg.get_pose_results())
num_target = len(pose_preds)
# Show mode
img = frame_msg.get_image()
canvas = img.copy()
if self.mode == Mode.SHUFFLE:
mode_txt = 'Shuffle'
else:
mode_txt = 'Clone'
cv2.putText(canvas, mode_txt, (10, 50), cv2.FONT_HERSHEY_DUPLEX, 0.8,
(255, 126, 0), 1)
# Skip if target number is less than 2
if num_target >= 2:
# Generate new mapping if target number changes
if num_target != len(self.history['target_map']):
if self.mode == Mode.SHUFFLE:
self.history['target_map'] = self._get_swap_map(num_target)
else:
self.history['target_map'] = np.repeat(
np.random.choice(num_target), num_target)
# # Draw on canvas
for tar_idx, src_idx in enumerate(self.history['target_map']):
face_src = self._get_face_info(pose_preds[src_idx])
face_tar = self._get_face_info(pose_preds[tar_idx])
canvas = self._swap_face(img, canvas, face_src, face_tar)
return canvas
def _crop_face_by_contour(self, img, contour):
mask = np.zeros(img.shape[:2], dtype=np.uint8)
cv2.fillPoly(mask, [contour.astype(np.int32)], 1)
mask = cv2.dilate(
mask, kernel=np.ones((9, 9), dtype=np.uint8), anchor=(4, 0))
x1, y1, w, h = cv2.boundingRect(mask)
x2 = x1 + w
y2 = y1 + h
bbox = np.array([x1, y1, x2, y2], dtype=np.int64)
patch = img[y1:y2, x1:x2]
mask = mask[y1:y2, x1:x2]
return bbox, patch, mask
def _swap_face(self, img_src, img_tar, face_src, face_tar):
if face_src['dataset'] == face_tar['dataset']:
# Use full keypoints for face alignment
kpts_src = face_src['contour']
kpts_tar = face_tar['contour']
else:
# Use only common landmarks (eyes and nose) for face alignment if
# source and target have differenet data type
# (e.g. human vs animal)
kpts_src = face_src['landmarks']
kpts_tar = face_tar['landmarks']
# Get everything local
bbox_src, patch_src, mask_src = self._crop_face_by_contour(
img_src, face_src['contour'])
bbox_tar, _, mask_tar = self._crop_face_by_contour(
img_tar, face_tar['contour'])
kpts_src = kpts_src - bbox_src[:2]
kpts_tar = kpts_tar - bbox_tar[:2]
# Compute affine transformation matrix
trans_mat, _ = cv2.estimateAffine2D(
kpts_src.astype(np.float32), kpts_tar.astype(np.float32))
patch_warp = cv2.warpAffine(
patch_src,
trans_mat,
dsize=tuple(bbox_tar[2:] - bbox_tar[:2]),
borderValue=(0, 0, 0))
mask_warp = cv2.warpAffine(
mask_src,
trans_mat,
dsize=tuple(bbox_tar[2:] - bbox_tar[:2]),
borderValue=(0, 0, 0))
# Target mask
mask_tar = mask_tar & mask_warp
mask_tar_soft = cv2.GaussianBlur(mask_tar * 255, (3, 3), 3)
# Blending
center = tuple((0.5 * (bbox_tar[:2] + bbox_tar[2:])).astype(np.int64))
img_tar = cv2.seamlessClone(patch_warp, img_tar, mask_tar_soft, center,
cv2.NORMAL_CLONE)
return img_tar
@staticmethod
def _get_face_info(pose_pred):
keypoints = pose_pred['keypoints'][:, :2]
model_cfg = pose_pred['model_cfg']
dataset_info = DatasetInfo(model_cfg.data.test.dataset_info)
face_info = {
'dataset': dataset_info.dataset_name,
'landmarks': None, # For alignment
'contour': None, # For mask generation
'bbox': None # For image warping
}
# Fall back to hard coded keypoint id
if face_info['dataset'] == 'coco':
face_info['landmarks'] = np.stack([
keypoints[1], # left eye
keypoints[2], # right eye
keypoints[0], # nose
0.5 * (keypoints[5] + keypoints[6]), # neck (shoulder center)
])
elif face_info['dataset'] == 'coco_wholebody':
face_info['landmarks'] = np.stack([
keypoints[1], # left eye
keypoints[2], # right eye
keypoints[0], # nose
keypoints[32], # chin
])
contour_ids = list(range(23, 40)) + list(range(40, 50))[::-1]
face_info['contour'] = keypoints[contour_ids]
elif face_info['dataset'] == 'ap10k':
face_info['landmarks'] = np.stack([
keypoints[0], # left eye
keypoints[1], # right eye
keypoints[2], # nose
keypoints[3], # neck
])
elif face_info['dataset'] == 'animalpose':
face_info['landmarks'] = np.stack([
keypoints[0], # left eye
keypoints[1], # right eye
keypoints[4], # nose
keypoints[5], # throat
])
elif face_info['dataset'] == 'wflw':
face_info['landmarks'] = np.stack([
keypoints[97], # left eye
keypoints[96], # right eye
keypoints[54], # nose
keypoints[16], # chine
])
contour_ids = list(range(33))[::-1] + list(range(33, 38)) + list(
range(42, 47))
face_info['contour'] = keypoints[contour_ids]
else:
raise ValueError('Can not obtain face landmark information'
f'from dataset: {face_info["type"]}')
# Face region
if face_info['contour'] is None:
# Manually defined counter of face region
left_eye, right_eye, nose = face_info['landmarks'][:3]
eye_center = 0.5 * (left_eye + right_eye)
w_vec = right_eye - left_eye
eye_dist = np.linalg.norm(w_vec) + 1e-6
w_vec = w_vec / eye_dist
h_vec = np.array([w_vec[1], -w_vec[0]], dtype=w_vec.dtype)
w = max(0.5 * eye_dist, np.abs(np.dot(nose - eye_center, w_vec)))
h = np.abs(np.dot(nose - eye_center, h_vec))
left_top = eye_center + 1.5 * w * w_vec - 0.5 * h * h_vec
right_top = eye_center - 1.5 * w * w_vec - 0.5 * h * h_vec
left_bottom = eye_center + 1.5 * w * w_vec + 4 * h * h_vec
right_bottom = eye_center - 1.5 * w * w_vec + 4 * h * h_vec
face_info['contour'] = np.stack(
[left_top, right_top, right_bottom, left_bottom])
# Get tight bbox of face region
face_info['bbox'] = np.array([
face_info['contour'][:, 0].min(), face_info['contour'][:, 1].min(),
face_info['contour'][:, 0].max(), face_info['contour'][:, 1].max()
]).astype(np.int64)
return face_info
@staticmethod
def _merge_pose_results(pose_results):
preds = []
if pose_results is not None:
for prefix, pose_result in enumerate(pose_results):
model_cfg = pose_result['model_cfg']
for idx, _pred in enumerate(pose_result['preds']):
pred = _pred.copy()
pred['id'] = f'{prefix}.{_pred.get("track_id", str(idx))}'
pred['model_cfg'] = model_cfg
preds.append(pred)
return preds
@staticmethod
def _get_swap_map(num_target):
ids = np.random.choice(num_target, num_target, replace=False)
target_map = ids[(ids + 1) % num_target]
return target_map
@@ -0,0 +1,65 @@
# Copyright (c) OpenMMLab. All rights reserved.
from abc import abstractmethod
from typing import Dict, List, Optional, Union
import numpy as np
from ..utils import FrameMessage, Message
from .node import Node
class FrameDrawingNode(Node):
"""Base class for Node that draw on single frame images.
Args:
name (str, optional): The node name (also thread name).
frame_buffer (str): The name of the input buffer.
output_buffer (str | list): The name(s) of the output buffer(s).
enable_key (str | int, optional): Set a hot-key to toggle
enable/disable of the node. If an int value is given, it will be
treated as an ascii code of a key. Please note:
1. If enable_key is set, the bypass method need to be
overridden to define the node behavior when disabled
2. Some hot-key has been use for particular use. For example:
'q', 'Q' and 27 are used for quit
Default: None
enable (bool): Default enable/disable status. Default: True.
"""
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True):
super().__init__(name=name, enable_key=enable_key)
# Register buffers
self.register_input_buffer(frame_buffer, 'frame', essential=True)
self.register_output_buffer(output_buffer)
self._enabled = enable
def process(self, input_msgs: Dict[str, Message]) -> Union[Message, None]:
frame_msg = input_msgs['frame']
img = self.draw(frame_msg)
frame_msg.set_image(img)
return frame_msg
def bypass(self, input_msgs: Dict[str, Message]) -> Union[Message, None]:
return input_msgs['frame']
@abstractmethod
def draw(self, frame_msg: FrameMessage) -> np.ndarray:
"""Draw on the frame image with information from the single frame.
Args:
frame_meg (FrameMessage): The frame to get information from and
draw on.
Returns:
array: The output image
"""
@@ -0,0 +1,917 @@
# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, List, Optional, Tuple, Union
import cv2
import numpy as np
from mmcv import color_val
from mmpose.core import (apply_bugeye_effect, apply_sunglasses_effect,
imshow_bboxes, imshow_keypoints)
from mmpose.datasets import DatasetInfo
from ..utils import (FrameMessage, copy_and_paste, expand_and_clamp,
get_cached_file_path, get_eye_keypoint_ids,
get_face_keypoint_ids, get_wrist_keypoint_ids,
load_image_from_disk_or_url, screen_matting)
from .builder import NODES
from .frame_drawing_node import FrameDrawingNode
try:
import psutil
psutil_proc = psutil.Process()
except (ImportError, ModuleNotFoundError):
psutil_proc = None
@NODES.register_module()
class PoseVisualizerNode(FrameDrawingNode):
"""Draw the bbox and keypoint detection results.
Args:
name (str, optional): The node name (also thread name).
frame_buffer (str): The name of the input buffer.
output_buffer (str|list): The name(s) of the output buffer(s).
enable_key (str|int, optional): Set a hot-key to toggle enable/disable
of the node. If an int value is given, it will be treated as an
ascii code of a key. Please note:
1. If enable_key is set, the bypass method need to be
overridden to define the node behavior when disabled
2. Some hot-key has been use for particular use. For example:
'q', 'Q' and 27 are used for quit
Default: None
enable (bool): Default enable/disable status. Default: True.
kpt_thr (float): The threshold of keypoint score. Default: 0.3.
radius (int): The radius of keypoint. Default: 4.
thickness (int): The thickness of skeleton. Default: 2.
bbox_color (str|tuple|dict): If a single color (a str like 'green' or
a tuple like (0, 255, 0)), it will used to draw the bbox.
Optionally, a dict can be given as a map from class labels to
colors.
"""
default_bbox_color = {
'person': (148, 139, 255),
'cat': (255, 255, 0),
'dog': (255, 255, 0),
}
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
kpt_thr: float = 0.3,
radius: int = 4,
thickness: int = 2,
bbox_color: Optional[Union[str, Tuple, Dict]] = None):
super().__init__(name, frame_buffer, output_buffer, enable_key, enable)
self.kpt_thr = kpt_thr
self.radius = radius
self.thickness = thickness
if bbox_color is None:
self.bbox_color = self.default_bbox_color
elif isinstance(bbox_color, dict):
self.bbox_color = {k: color_val(v) for k, v in bbox_color.items()}
else:
self.bbox_color = color_val(bbox_color)
def draw(self, frame_msg):
canvas = frame_msg.get_image()
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
for pose_result in frame_msg.get_pose_results():
model_cfg = pose_result['model_cfg']
dataset_info = DatasetInfo(model_cfg.dataset_info)
# Extract bboxes and poses
bbox_preds = []
bbox_labels = []
pose_preds = []
for pred in pose_result['preds']:
if 'bbox' in pred:
bbox_preds.append(pred['bbox'])
bbox_labels.append(pred.get('label', None))
pose_preds.append(pred['keypoints'])
# Get bbox colors
if isinstance(self.bbox_color, dict):
bbox_colors = [
self.bbox_color.get(label, (0, 255, 0))
for label in bbox_labels
]
else:
bbox_labels = self.bbox_color
# Draw bboxes
if bbox_preds:
bboxes = np.vstack(bbox_preds)
imshow_bboxes(
canvas,
bboxes,
labels=bbox_labels,
colors=bbox_colors,
text_color='white',
font_scale=0.5,
show=False)
# Draw poses
if pose_preds:
imshow_keypoints(
canvas,
pose_preds,
skeleton=dataset_info.skeleton,
kpt_score_thr=0.3,
pose_kpt_color=dataset_info.pose_kpt_color,
pose_link_color=dataset_info.pose_link_color,
radius=self.radius,
thickness=self.thickness)
return canvas
@NODES.register_module()
class SunglassesNode(FrameDrawingNode):
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
src_img_path: Optional[str] = None):
super().__init__(name, frame_buffer, output_buffer, enable_key, enable)
if src_img_path is None:
# The image attributes to:
# https://www.vecteezy.com/free-vector/glass
# Glass Vectors by Vecteezy
src_img_path = 'demo/resources/sunglasses.jpg'
self.src_img = load_image_from_disk_or_url(src_img_path)
def draw(self, frame_msg):
canvas = frame_msg.get_image()
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
for pose_result in pose_results:
model_cfg = pose_result['model_cfg']
preds = pose_result['preds']
left_eye_idx, right_eye_idx = get_eye_keypoint_ids(model_cfg)
canvas = apply_sunglasses_effect(canvas, preds, self.src_img,
left_eye_idx, right_eye_idx)
return canvas
@NODES.register_module()
class SpriteNode(FrameDrawingNode):
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
src_img_path: Optional[str] = None):
super().__init__(name, frame_buffer, output_buffer, enable_key, enable)
if src_img_path is None:
# Sprites of Touhou characters :)
# Come from https://www.deviantart.com/shadowbendy/art/Touhou-rpg-maker-vx-Sprite-1-812746920 # noqa: E501
src_img_path = (
'https://user-images.githubusercontent.com/'
'26739999/151532276-33f968d9-917f-45e3-8a99-ebde60be83bb.png')
self.src_img = load_image_from_disk_or_url(
src_img_path, cv2.IMREAD_UNCHANGED)[:144, :108]
tmp = np.array(np.split(self.src_img, range(36, 144, 36), axis=0))
tmp = np.array(np.split(tmp, range(36, 108, 36), axis=2))
self.sprites = tmp
self.pos = None
self.anime_frame = 0
def apply_sprite_effect(self,
img,
pose_results,
left_hand_index,
right_hand_index,
kpt_thr=0.5):
"""Apply sprite effect.
Args:
img (np.ndarray): Image data.
pose_results (list[dict]): The pose estimation results containing:
- "keypoints" ([K,3]): detection result in [x, y, score]
left_hand_index (int): Keypoint index of left hand
right_hand_index (int): Keypoint index of right hand
kpt_thr (float): The score threshold of required keypoints.
"""
hm, wm = self.sprites.shape[2:4]
# anchor points in the sunglasses mask
if self.pos is None:
self.pos = [img.shape[0] // 2, img.shape[1] // 2]
if len(pose_results) == 0:
return img
kpts = pose_results[0]['keypoints']
if kpts[left_hand_index, 2] < kpt_thr and kpts[right_hand_index,
2] < kpt_thr:
aim = self.pos
else:
kpt_lhand = kpts[left_hand_index, :2][::-1]
kpt_rhand = kpts[right_hand_index, :2][::-1]
def distance(a, b):
return (a[0] - b[0])**2 + (a[1] - b[1])**2
# Go to the nearest hand
if distance(kpt_lhand, self.pos) < distance(kpt_rhand, self.pos):
aim = kpt_lhand
else:
aim = kpt_rhand
pos_thr = 15
if aim[0] < self.pos[0] - pos_thr:
# Go down
sprite = self.sprites[self.anime_frame][3]
self.pos[0] -= 1
elif aim[0] > self.pos[0] + pos_thr:
# Go up
sprite = self.sprites[self.anime_frame][0]
self.pos[0] += 1
elif aim[1] < self.pos[1] - pos_thr:
# Go right
sprite = self.sprites[self.anime_frame][1]
self.pos[1] -= 1
elif aim[1] > self.pos[1] + pos_thr:
# Go left
sprite = self.sprites[self.anime_frame][2]
self.pos[1] += 1
else:
# Stay
self.anime_frame = 0
sprite = self.sprites[self.anime_frame][0]
if self.anime_frame < 2:
self.anime_frame += 1
else:
self.anime_frame = 0
x = self.pos[0] - hm // 2
y = self.pos[1] - wm // 2
x = max(0, min(x, img.shape[0] - hm))
y = max(0, min(y, img.shape[0] - wm))
# Overlay image with transparent
img[x:x + hm, y:y +
wm] = (img[x:x + hm, y:y + wm] * (1 - sprite[:, :, 3:] / 255) +
sprite[:, :, :3] * (sprite[:, :, 3:] / 255)).astype('uint8')
return img
def draw(self, frame_msg):
canvas = frame_msg.get_image()
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
for pose_result in pose_results:
model_cfg = pose_result['model_cfg']
preds = pose_result['preds']
# left_hand_idx, right_hand_idx = get_wrist_keypoint_ids(model_cfg) # noqa: E501
left_hand_idx, right_hand_idx = get_eye_keypoint_ids(model_cfg)
canvas = self.apply_sprite_effect(canvas, preds, left_hand_idx,
right_hand_idx)
return canvas
@NODES.register_module()
class BackgroundNode(FrameDrawingNode):
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
src_img_path: Optional[str] = None,
cls_ids: Optional[List] = None,
cls_names: Optional[List] = None):
super().__init__(name, frame_buffer, output_buffer, enable_key, enable)
self.cls_ids = cls_ids
self.cls_names = cls_names
if src_img_path is None:
src_img_path = 'https://user-images.githubusercontent.com/'\
'11788150/149731957-abd5c908-9c7f-45b2-b7bf-'\
'821ab30c6a3e.jpg'
self.src_img = load_image_from_disk_or_url(src_img_path)
def apply_background_effect(self,
img,
det_results,
background_img,
effect_region=(0.2, 0.2, 0.8, 0.8)):
"""Change background.
Args:
img (np.ndarray): Image data.
det_results (list[dict]): The detection results containing:
- "cls_id" (int): Class index.
- "label" (str): Class label (e.g. 'person').
- "bbox" (ndarray:(5, )): bounding box result
[x, y, w, h, score].
- "mask" (ndarray:(w, h)): instance segmentation result.
background_img (np.ndarray): Background image.
effect_region (tuple(4, )): The region to apply mask,
the coordinates are normalized (x1, y1, x2, y2).
"""
if len(det_results) > 0:
# Choose the one with the highest score.
det_result = det_results[0]
bbox = det_result['bbox']
mask = det_result['mask'].astype(np.uint8)
img = copy_and_paste(img, background_img, mask, bbox,
effect_region)
return img
else:
return background_img
def draw(self, frame_msg):
canvas = frame_msg.get_image()
if canvas.shape != self.src_img.shape:
self.src_img = cv2.resize(self.src_img, canvas.shape[:2])
det_results = frame_msg.get_detection_results()
if not det_results:
return canvas
full_preds = []
for det_result in det_results:
preds = det_result['preds']
if self.cls_ids:
# Filter results by class ID
filtered_preds = [
p for p in preds if p['cls_id'] in self.cls_ids
]
elif self.cls_names:
# Filter results by class name
filtered_preds = [
p for p in preds if p['label'] in self.cls_names
]
else:
filtered_preds = preds
full_preds.extend(filtered_preds)
canvas = self.apply_background_effect(canvas, full_preds, self.src_img)
return canvas
@NODES.register_module()
class SaiyanNode(FrameDrawingNode):
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
hair_img_path: Optional[str] = None,
light_video_path: Optional[str] = None,
cls_ids: Optional[List] = None,
cls_names: Optional[List] = None):
super().__init__(name, frame_buffer, output_buffer, enable_key, enable)
self.cls_ids = cls_ids
self.cls_names = cls_names
if hair_img_path is None:
hair_img_path = 'https://user-images.githubusercontent.com/'\
'11788150/149732117-fcd2d804-dc2c-426c-bee7-'\
'94be6146e05c.png'
self.hair_img = load_image_from_disk_or_url(hair_img_path)
if light_video_path is None:
light_video_path = get_cached_file_path(
'https://'
'user-images.githubusercontent.com/11788150/149732080'
'-ea6cfeda-0dc5-4bbb-892a-3831e5580520.mp4')
self.light_video_path = light_video_path
self.light_video = cv2.VideoCapture(self.light_video_path)
def apply_saiyan_effect(self,
img,
pose_results,
saiyan_img,
light_frame,
face_indices,
bbox_thr=0.3,
kpt_thr=0.5):
"""Apply saiyan hair effect.
Args:
img (np.ndarray): Image data.
pose_results (list[dict]): The pose estimation results containing:
- "keypoints" ([K,3]): keypoint detection result
in [x, y, score]
saiyan_img (np.ndarray): Saiyan image with transparent background.
light_frame (np.ndarray): Light image with green screen.
face_indices (int): Keypoint index of the face
kpt_thr (float): The score threshold of required keypoints.
"""
img = img.copy()
im_shape = img.shape
# Apply lightning effects.
light_mask = screen_matting(light_frame, color='green')
# anchor points in the mask
pts_src = np.array(
[
[84, 398], # face kpt 0
[331, 393], # face kpt 16
[84, 145],
[331, 140]
],
dtype=np.float32)
for pose in pose_results:
bbox = pose['bbox']
if bbox[-1] < bbox_thr:
continue
mask_inst = pose['mask']
# cache
fg = img[np.where(mask_inst)]
bbox = expand_and_clamp(bbox[:4], im_shape, s=3.0)
# Apply light effects between fg and bg
img = copy_and_paste(
light_frame,
img,
light_mask,
effect_region=(bbox[0] / im_shape[1], bbox[1] / im_shape[0],
bbox[2] / im_shape[1], bbox[3] / im_shape[0]))
# pop
img[np.where(mask_inst)] = fg
# Apply Saiyan hair effects
kpts = pose['keypoints']
if kpts[face_indices[0], 2] < kpt_thr or kpts[face_indices[16],
2] < kpt_thr:
continue
kpt_0 = kpts[face_indices[0], :2]
kpt_16 = kpts[face_indices[16], :2]
# orthogonal vector
vo = (kpt_0 - kpt_16)[::-1] * [-1, 1]
# anchor points in the image by eye positions
pts_tar = np.vstack([kpt_0, kpt_16, kpt_0 + vo, kpt_16 + vo])
h_mat, _ = cv2.findHomography(pts_src, pts_tar)
patch = cv2.warpPerspective(
saiyan_img,
h_mat,
dsize=(img.shape[1], img.shape[0]),
borderValue=(0, 0, 0))
mask_patch = cv2.cvtColor(patch, cv2.COLOR_BGR2GRAY)
mask_patch = (mask_patch > 1).astype(np.uint8)
img = cv2.copyTo(patch, mask_patch, img)
return img
def draw(self, frame_msg):
canvas = frame_msg.get_image()
det_results = frame_msg.get_detection_results()
if not det_results:
return canvas
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
for pose_result in pose_results:
model_cfg = pose_result['model_cfg']
preds = pose_result['preds']
face_indices = get_face_keypoint_ids(model_cfg)
ret, frame = self.light_video.read()
if not ret:
self.light_video = cv2.VideoCapture(self.light_video_path)
ret, frame = self.light_video.read()
canvas = self.apply_saiyan_effect(canvas, preds, self.hair_img,
frame, face_indices)
return canvas
@NODES.register_module()
class MoustacheNode(FrameDrawingNode):
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
src_img_path: Optional[str] = None):
super().__init__(name, frame_buffer, output_buffer, enable_key, enable)
if src_img_path is None:
src_img_path = 'https://user-images.githubusercontent.com/'\
'11788150/149732141-3afbab55-252a-428c-b6d8'\
'-0e352f432651.jpeg'
self.src_img = load_image_from_disk_or_url(src_img_path)
def apply_moustache_effect(self,
img,
pose_results,
moustache_img,
face_indices,
kpt_thr=0.5):
"""Apply moustache effect.
Args:
img (np.ndarray): Image data.
pose_results (list[dict]): The pose estimation results containing:
- "keypoints" ([K,3]): keypoint detection result
in [x, y, score]
moustache_img (np.ndarray): Moustache image with white background.
left_eye_index (int): Keypoint index of left eye
right_eye_index (int): Keypoint index of right eye
kpt_thr (float): The score threshold of required keypoints.
"""
hm, wm = moustache_img.shape[:2]
# anchor points in the moustache mask
pts_src = np.array([[1164, 741], [1729, 741], [1164, 1244],
[1729, 1244]],
dtype=np.float32)
for pose in pose_results:
kpts = pose['keypoints']
if kpts[face_indices[32], 2] < kpt_thr \
or kpts[face_indices[34], 2] < kpt_thr \
or kpts[face_indices[61], 2] < kpt_thr \
or kpts[face_indices[63], 2] < kpt_thr:
continue
kpt_32 = kpts[face_indices[32], :2]
kpt_34 = kpts[face_indices[34], :2]
kpt_61 = kpts[face_indices[61], :2]
kpt_63 = kpts[face_indices[63], :2]
# anchor points in the image by eye positions
pts_tar = np.vstack([kpt_32, kpt_34, kpt_61, kpt_63])
h_mat, _ = cv2.findHomography(pts_src, pts_tar)
patch = cv2.warpPerspective(
moustache_img,
h_mat,
dsize=(img.shape[1], img.shape[0]),
borderValue=(255, 255, 255))
# mask the white background area in the patch with a threshold 200
mask = cv2.cvtColor(patch, cv2.COLOR_BGR2GRAY)
mask = (mask < 200).astype(np.uint8)
img = cv2.copyTo(patch, mask, img)
return img
def draw(self, frame_msg):
canvas = frame_msg.get_image()
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
for pose_result in pose_results:
model_cfg = pose_result['model_cfg']
preds = pose_result['preds']
face_indices = get_face_keypoint_ids(model_cfg)
canvas = self.apply_moustache_effect(canvas, preds, self.src_img,
face_indices)
return canvas
@NODES.register_module()
class BugEyeNode(FrameDrawingNode):
def draw(self, frame_msg):
canvas = frame_msg.get_image()
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
for pose_result in pose_results:
model_cfg = pose_result['model_cfg']
preds = pose_result['preds']
left_eye_idx, right_eye_idx = get_eye_keypoint_ids(model_cfg)
canvas = apply_bugeye_effect(canvas, preds, left_eye_idx,
right_eye_idx)
return canvas
@NODES.register_module()
class NoticeBoardNode(FrameDrawingNode):
default_content_lines = ['This is a notice board!']
def __init__(
self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
content_lines: Optional[List[str]] = None,
x_offset: int = 20,
y_offset: int = 20,
y_delta: int = 15,
text_color: Union[str, Tuple[int, int, int]] = 'black',
background_color: Union[str, Tuple[int, int, int]] = (255, 183, 0),
text_scale: float = 0.4,
):
super().__init__(name, frame_buffer, output_buffer, enable_key, enable)
self.x_offset = x_offset
self.y_offset = y_offset
self.y_delta = y_delta
self.text_color = color_val(text_color)
self.background_color = color_val(background_color)
self.text_scale = text_scale
if content_lines:
self.content_lines = content_lines
else:
self.content_lines = self.default_content_lines
def draw(self, frame_msg: FrameMessage) -> np.ndarray:
img = frame_msg.get_image()
canvas = np.full(img.shape, self.background_color, dtype=img.dtype)
x = self.x_offset
y = self.y_offset
max_len = max([len(line) for line in self.content_lines])
def _put_line(line=''):
nonlocal y
cv2.putText(canvas, line, (x, y), cv2.FONT_HERSHEY_DUPLEX,
self.text_scale, self.text_color, 1)
y += self.y_delta
for line in self.content_lines:
_put_line(line)
x1 = max(0, self.x_offset)
x2 = min(img.shape[1], int(x + max_len * self.text_scale * 20))
y1 = max(0, self.y_offset - self.y_delta)
y2 = min(img.shape[0], y)
src1 = canvas[y1:y2, x1:x2]
src2 = img[y1:y2, x1:x2]
img[y1:y2, x1:x2] = cv2.addWeighted(src1, 0.5, src2, 0.5, 0)
return img
@NODES.register_module()
class HatNode(FrameDrawingNode):
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
src_img_path: Optional[str] = None):
super().__init__(name, frame_buffer, output_buffer, enable_key)
if src_img_path is None:
# The image attributes to:
# http://616pic.com/sucai/1m9i70p52.html
src_img_path = 'https://user-images.githubusercontent.' \
'com/28900607/149766271-2f591c19-9b67-4' \
'd92-8f94-c272396ca141.png'
self.src_img = load_image_from_disk_or_url(src_img_path,
cv2.IMREAD_UNCHANGED)
@staticmethod
def apply_hat_effect(img,
pose_results,
hat_img,
left_eye_index,
right_eye_index,
kpt_thr=0.5):
"""Apply hat effect.
Args:
img (np.ndarray): Image data.
pose_results (list[dict]): The pose estimation results containing:
- "keypoints" ([K,3]): keypoint detection result in
[x, y, score]
hat_img (np.ndarray): Hat image with white alpha channel.
left_eye_index (int): Keypoint index of left eye
right_eye_index (int): Keypoint index of right eye
kpt_thr (float): The score threshold of required keypoints.
"""
img_orig = img.copy()
img = img_orig.copy()
hm, wm = hat_img.shape[:2]
# anchor points in the sunglasses mask
a = 0.3
b = 0.7
pts_src = np.array([[a * wm, a * hm], [a * wm, b * hm],
[b * wm, a * hm], [b * wm, b * hm]],
dtype=np.float32)
for pose in pose_results:
kpts = pose['keypoints']
if kpts[left_eye_index, 2] < kpt_thr or \
kpts[right_eye_index, 2] < kpt_thr:
continue
kpt_leye = kpts[left_eye_index, :2]
kpt_reye = kpts[right_eye_index, :2]
# orthogonal vector to the left-to-right eyes
vo = 0.5 * (kpt_reye - kpt_leye)[::-1] * [-1, 1]
veye = 0.5 * (kpt_reye - kpt_leye)
# anchor points in the image by eye positions
pts_tar = np.vstack([
kpt_reye + 1 * veye + 5 * vo, kpt_reye + 1 * veye + 1 * vo,
kpt_leye - 1 * veye + 5 * vo, kpt_leye - 1 * veye + 1 * vo
])
h_mat, _ = cv2.findHomography(pts_src, pts_tar)
patch = cv2.warpPerspective(
hat_img,
h_mat,
dsize=(img.shape[1], img.shape[0]),
borderValue=(255, 255, 255))
# mask the white background area in the patch with a threshold 200
mask = (patch[:, :, -1] > 128)
patch = patch[:, :, :-1]
mask = mask * (cv2.cvtColor(patch, cv2.COLOR_BGR2GRAY) > 30)
mask = mask.astype(np.uint8)
img = cv2.copyTo(patch, mask, img)
return img
def draw(self, frame_msg):
canvas = frame_msg.get_image()
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
for pose_result in pose_results:
model_cfg = pose_result['model_cfg']
preds = pose_result['preds']
left_eye_idx, right_eye_idx = get_eye_keypoint_ids(model_cfg)
canvas = self.apply_hat_effect(canvas, preds, self.src_img,
left_eye_idx, right_eye_idx)
return canvas
@NODES.register_module()
class FirecrackerNode(FrameDrawingNode):
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
src_img_path: Optional[str] = None):
super().__init__(name, frame_buffer, output_buffer, enable_key)
if src_img_path is None:
self.src_img_path = 'https://user-images.githubusercontent' \
'.com/28900607/149766281-6376055c-ed8b' \
'-472b-991f-60e6ae6ee1da.gif'
src_img = cv2.VideoCapture(self.src_img_path)
self.frame_list = []
ret, frame = src_img.read()
while frame is not None:
self.frame_list.append(frame)
ret, frame = src_img.read()
self.num_frames = len(self.frame_list)
self.frame_idx = 0
self.frame_period = 4 # each frame in gif lasts for 4 frames in video
@staticmethod
def apply_firecracker_effect(img,
pose_results,
firecracker_img,
left_wrist_idx,
right_wrist_idx,
kpt_thr=0.5):
"""Apply firecracker effect.
Args:
img (np.ndarray): Image data.
pose_results (list[dict]): The pose estimation results containing:
- "keypoints" ([K,3]): keypoint detection result in
[x, y, score]
firecracker_img (np.ndarray): Firecracker image with white
background.
left_wrist_idx (int): Keypoint index of left wrist
right_wrist_idx (int): Keypoint index of right wrist
kpt_thr (float): The score threshold of required keypoints.
"""
hm, wm = firecracker_img.shape[:2]
# anchor points in the firecracker mask
pts_src = np.array([[0. * wm, 0. * hm], [0. * wm, 1. * hm],
[1. * wm, 0. * hm], [1. * wm, 1. * hm]],
dtype=np.float32)
h, w = img.shape[:2]
h_tar = h / 3
w_tar = h_tar / hm * wm
for pose in pose_results:
kpts = pose['keypoints']
if kpts[left_wrist_idx, 2] > kpt_thr:
kpt_lwrist = kpts[left_wrist_idx, :2]
# anchor points in the image by eye positions
pts_tar = np.vstack([
kpt_lwrist - [w_tar / 2, 0],
kpt_lwrist - [w_tar / 2, -h_tar],
kpt_lwrist + [w_tar / 2, 0],
kpt_lwrist + [w_tar / 2, h_tar]
])
h_mat, _ = cv2.findHomography(pts_src, pts_tar)
patch = cv2.warpPerspective(
firecracker_img,
h_mat,
dsize=(img.shape[1], img.shape[0]),
borderValue=(255, 255, 255))
# mask the white background area in the patch with
# a threshold 200
mask = cv2.cvtColor(patch, cv2.COLOR_BGR2GRAY)
mask = (mask < 240).astype(np.uint8)
img = cv2.copyTo(patch, mask, img)
if kpts[right_wrist_idx, 2] > kpt_thr:
kpt_rwrist = kpts[right_wrist_idx, :2]
# anchor points in the image by eye positions
pts_tar = np.vstack([
kpt_rwrist - [w_tar / 2, 0],
kpt_rwrist - [w_tar / 2, -h_tar],
kpt_rwrist + [w_tar / 2, 0],
kpt_rwrist + [w_tar / 2, h_tar]
])
h_mat, _ = cv2.findHomography(pts_src, pts_tar)
patch = cv2.warpPerspective(
firecracker_img,
h_mat,
dsize=(img.shape[1], img.shape[0]),
borderValue=(255, 255, 255))
# mask the white background area in the patch with
# a threshold 200
mask = cv2.cvtColor(patch, cv2.COLOR_BGR2GRAY)
mask = (mask < 240).astype(np.uint8)
img = cv2.copyTo(patch, mask, img)
return img
def draw(self, frame_msg):
canvas = frame_msg.get_image()
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
frame = self.frame_list[self.frame_idx // self.frame_period]
for pose_result in pose_results:
model_cfg = pose_result['model_cfg']
preds = pose_result['preds']
left_wrist_idx, right_wrist_idx = get_wrist_keypoint_ids(model_cfg)
canvas = self.apply_firecracker_effect(canvas, preds, frame,
left_wrist_idx,
right_wrist_idx)
self.frame_idx = (self.frame_idx + 1) % (
self.num_frames * self.frame_period)
return canvas
@@ -0,0 +1,296 @@
# Copyright (c) OpenMMLab. All rights reserved.
import logging
import time
from queue import Full, Queue
from threading import Thread
from typing import List, Optional, Union
import cv2
import numpy as np
from mmcv import color_val
from mmpose.utils.timer import RunningAverage
from .builder import NODES
from .node import Node
try:
import psutil
psutil_proc = psutil.Process()
except (ImportError, ModuleNotFoundError):
psutil_proc = None
@NODES.register_module()
class ModelResultBindingNode(Node):
def __init__(self, name: str, frame_buffer: str, result_buffer: str,
output_buffer: Union[str, List[str]]):
super().__init__(name=name, enable=True)
self.synchronous = None
# Cache the latest model result
self.last_result_msg = None
self.last_output_msg = None
# Inference speed analysis
self.frame_fps = RunningAverage(window=10)
self.frame_lag = RunningAverage(window=10)
self.result_fps = RunningAverage(window=10)
self.result_lag = RunningAverage(window=10)
# Register buffers
# Note that essential buffers will be set in set_runner() because
# it depends on the runner.synchronous attribute.
self.register_input_buffer(result_buffer, 'result', essential=False)
self.register_input_buffer(frame_buffer, 'frame', essential=False)
self.register_output_buffer(output_buffer)
def set_runner(self, runner):
super().set_runner(runner)
# Set synchronous according to the runner
if runner.synchronous:
self.synchronous = True
essential_input = 'result'
else:
self.synchronous = False
essential_input = 'frame'
# Set essential input buffer according to the synchronous setting
for buffer_info in self._input_buffers:
if buffer_info.input_name == essential_input:
buffer_info.essential = True
def process(self, input_msgs):
result_msg = input_msgs['result']
# Update last result
if result_msg is not None:
# Update result FPS
if self.last_result_msg is not None:
self.result_fps.update(
1.0 /
(result_msg.timestamp - self.last_result_msg.timestamp))
# Update inference latency
self.result_lag.update(time.time() - result_msg.timestamp)
# Update last inference result
self.last_result_msg = result_msg
if not self.synchronous:
# Asynchronous mode: Bind the latest result with the current frame.
frame_msg = input_msgs['frame']
self.frame_lag.update(time.time() - frame_msg.timestamp)
# Bind result to frame
if self.last_result_msg is not None:
frame_msg.set_full_results(
self.last_result_msg.get_full_results())
frame_msg.merge_route_info(
self.last_result_msg.get_route_info())
output_msg = frame_msg
else:
# Synchronous mode: Directly output the frame that the model result
# was obtained from.
self.frame_lag.update(time.time() - result_msg.timestamp)
output_msg = result_msg
# Update frame fps and lag
if self.last_output_msg is not None:
self.frame_lag.update(time.time() - output_msg.timestamp)
self.frame_fps.update(
1.0 / (output_msg.timestamp - self.last_output_msg.timestamp))
self.last_output_msg = output_msg
return output_msg
def _get_node_info(self):
info = super()._get_node_info()
info['result_fps'] = self.result_fps.average()
info['result_lag (ms)'] = self.result_lag.average() * 1000
info['frame_fps'] = self.frame_fps.average()
info['frame_lag (ms)'] = self.frame_lag.average() * 1000
return info
@NODES.register_module()
class MonitorNode(Node):
_default_ignore_items = ['timestamp']
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = False,
x_offset=20,
y_offset=20,
y_delta=15,
text_color='black',
background_color=(255, 183, 0),
text_scale=0.4,
ignore_items: Optional[List[str]] = None):
super().__init__(name=name, enable_key=enable_key, enable=enable)
self.x_offset = x_offset
self.y_offset = y_offset
self.y_delta = y_delta
self.text_color = color_val(text_color)
self.background_color = color_val(background_color)
self.text_scale = text_scale
if ignore_items is None:
self.ignore_items = self._default_ignore_items
else:
self.ignore_items = ignore_items
self.register_input_buffer(frame_buffer, 'frame', essential=True)
self.register_output_buffer(output_buffer)
def process(self, input_msgs):
frame_msg = input_msgs['frame']
frame_msg.update_route_info(
node_name='System Info',
node_type='dummy',
info=self._get_system_info())
img = frame_msg.get_image()
route_info = frame_msg.get_route_info()
img = self._show_route_info(img, route_info)
frame_msg.set_image(img)
return frame_msg
def _get_system_info(self):
sys_info = {}
if psutil_proc is not None:
sys_info['CPU(%)'] = psutil_proc.cpu_percent()
sys_info['Memory(%)'] = psutil_proc.memory_percent()
return sys_info
def _show_route_info(self, img, route_info):
canvas = np.full(img.shape, self.background_color, dtype=img.dtype)
x = self.x_offset
y = self.y_offset
max_len = 0
def _put_line(line=''):
nonlocal y, max_len
cv2.putText(canvas, line, (x, y), cv2.FONT_HERSHEY_DUPLEX,
self.text_scale, self.text_color, 1)
y += self.y_delta
max_len = max(max_len, len(line))
for node_info in route_info:
title = f'{node_info["node"]}({node_info["node_type"]})'
_put_line(title)
for k, v in node_info['info'].items():
if k in self.ignore_items:
continue
if isinstance(v, float):
v = f'{v:.1f}'
_put_line(f' {k}: {v}')
x1 = max(0, self.x_offset)
x2 = min(img.shape[1], int(x + max_len * self.text_scale * 20))
y1 = max(0, self.y_offset - self.y_delta)
y2 = min(img.shape[0], y)
src1 = canvas[y1:y2, x1:x2]
src2 = img[y1:y2, x1:x2]
img[y1:y2, x1:x2] = cv2.addWeighted(src1, 0.5, src2, 0.5, 0)
return img
def bypass(self, input_msgs):
return input_msgs['frame']
@NODES.register_module()
class RecorderNode(Node):
"""Record the frames into a local file."""
def __init__(
self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
out_video_file: str,
out_video_fps: int = 30,
out_video_codec: str = 'mp4v',
buffer_size: int = 30,
):
super().__init__(name=name, enable_key=None, enable=True)
self.queue = Queue(maxsize=buffer_size)
self.out_video_file = out_video_file
self.out_video_fps = out_video_fps
self.out_video_codec = out_video_codec
self.vwriter = None
# Register buffers
self.register_input_buffer(frame_buffer, 'frame', essential=True)
self.register_output_buffer(output_buffer)
# Start a new thread to write frame
self.t_record = Thread(target=self._record, args=(), daemon=True)
self.t_record.start()
def process(self, input_msgs):
frame_msg = input_msgs['frame']
img = frame_msg.get_image() if frame_msg is not None else None
img_queued = False
while not img_queued:
try:
self.queue.put(img, timeout=1)
img_queued = True
logging.info(f'{self.name}: recorder received one frame!')
except Full:
logging.info(f'{self.name}: recorder jamed!')
return frame_msg
def _record(self):
while True:
img = self.queue.get()
if img is None:
break
if self.vwriter is None:
fourcc = cv2.VideoWriter_fourcc(*self.out_video_codec)
fps = self.out_video_fps
frame_size = (img.shape[1], img.shape[0])
self.vwriter = cv2.VideoWriter(self.out_video_file, fourcc,
fps, frame_size)
assert self.vwriter.isOpened()
self.vwriter.write(img)
logging.info('Video recorder released!')
if self.vwriter is not None:
self.vwriter.release()
def on_exit(self):
try:
# Try putting a None into the output queue so the self.vwriter will
# be released after all queue frames have been written to file.
self.queue.put(None, timeout=1)
self.t_record.join(timeout=1)
except Full:
pass
if self.t_record.is_alive():
# Force to release self.vwriter
logging.info('Video recorder forced release!')
if self.vwriter is not None:
self.vwriter.release()
@@ -0,0 +1,84 @@
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Optional, Union
from .builder import NODES
from .node import Node
try:
from mmdet.apis import inference_detector, init_detector
has_mmdet = True
except (ImportError, ModuleNotFoundError):
has_mmdet = False
@NODES.register_module()
class DetectorNode(Node):
def __init__(self,
name: str,
model_config: str,
model_checkpoint: str,
input_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
device: str = 'cuda:0'):
# Check mmdetection is installed
assert has_mmdet, 'Please install mmdet to run the demo.'
super().__init__(name=name, enable_key=enable_key, enable=True)
self.model_config = model_config
self.model_checkpoint = model_checkpoint
self.device = device.lower()
# Init model
self.model = init_detector(
self.model_config,
self.model_checkpoint,
device=self.device.lower())
# Register buffers
self.register_input_buffer(input_buffer, 'input', essential=True)
self.register_output_buffer(output_buffer)
def bypass(self, input_msgs):
return input_msgs['input']
def process(self, input_msgs):
input_msg = input_msgs['input']
img = input_msg.get_image()
preds = inference_detector(self.model, img)
det_result = self._post_process(preds)
input_msg.add_detection_result(det_result, tag=self.name)
return input_msg
def _post_process(self, preds):
if isinstance(preds, tuple):
dets = preds[0]
segms = preds[1]
else:
dets = preds
segms = [None] * len(dets)
assert len(dets) == len(self.model.CLASSES)
assert len(segms) == len(self.model.CLASSES)
result = {'preds': [], 'model_cfg': self.model.cfg.copy()}
for i, (cls_name, bboxes,
masks) in enumerate(zip(self.model.CLASSES, dets, segms)):
if masks is None:
masks = [None] * len(bboxes)
else:
assert len(masks) == len(bboxes)
preds_i = [{
'cls_id': i,
'label': cls_name,
'bbox': bbox,
'mask': mask,
} for (bbox, mask) in zip(bboxes, masks)]
result['preds'].extend(preds_i)
return result
@@ -0,0 +1,122 @@
# Copyright (c) OpenMMLab. All rights reserved.
import time
from typing import Dict, List, Optional, Union
from mmpose.apis import (get_track_id, inference_top_down_pose_model,
init_pose_model)
from ..utils import Message
from .builder import NODES
from .node import Node
@NODES.register_module()
class TopDownPoseEstimatorNode(Node):
def __init__(self,
name: str,
model_config: str,
model_checkpoint: str,
input_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
device: str = 'cuda:0',
cls_ids: Optional[List] = None,
cls_names: Optional[List] = None,
bbox_thr: float = 0.5):
super().__init__(name=name, enable_key=enable_key, enable=enable)
# Init model
self.model_config = model_config
self.model_checkpoint = model_checkpoint
self.device = device.lower()
self.cls_ids = cls_ids
self.cls_names = cls_names
self.bbox_thr = bbox_thr
# Init model
self.model = init_pose_model(
self.model_config,
self.model_checkpoint,
device=self.device.lower())
# Store history for pose tracking
self.track_info = {
'next_id': 0,
'last_pose_preds': [],
'last_time': None
}
# Register buffers
self.register_input_buffer(input_buffer, 'input', essential=True)
self.register_output_buffer(output_buffer)
def bypass(self, input_msgs):
return input_msgs['input']
def process(self, input_msgs: Dict[str, Message]) -> Message:
input_msg = input_msgs['input']
img = input_msg.get_image()
det_results = input_msg.get_detection_results()
if det_results is None:
raise ValueError(
'No detection results are found in the frame message.'
f'{self.__class__.__name__} should be used after a '
'detector node.')
full_det_preds = []
for det_result in det_results:
det_preds = det_result['preds']
if self.cls_ids:
# Filter detection results by class ID
det_preds = [
p for p in det_preds if p['cls_id'] in self.cls_ids
]
elif self.cls_names:
# Filter detection results by class name
det_preds = [
p for p in det_preds if p['label'] in self.cls_names
]
full_det_preds.extend(det_preds)
# Inference pose
pose_preds, _ = inference_top_down_pose_model(
self.model,
img,
full_det_preds,
bbox_thr=self.bbox_thr,
format='xyxy')
# Pose tracking
current_time = time.time()
if self.track_info['last_time'] is None:
fps = None
elif self.track_info['last_time'] >= current_time:
fps = None
else:
fps = 1.0 / (current_time - self.track_info['last_time'])
pose_preds, next_id = get_track_id(
pose_preds,
self.track_info['last_pose_preds'],
self.track_info['next_id'],
use_oks=False,
tracking_thr=0.3,
use_one_euro=True,
fps=fps)
self.track_info['next_id'] = next_id
self.track_info['last_pose_preds'] = pose_preds.copy()
self.track_info['last_time'] = current_time
pose_result = {
'preds': pose_preds,
'model_cfg': self.model.cfg.copy(),
}
input_msg.add_pose_result(pose_result, tag=self.name)
return input_msg
@@ -0,0 +1,372 @@
# Copyright (c) OpenMMLab. All rights reserved.
import logging
import time
from abc import ABCMeta, abstractmethod
from dataclasses import dataclass
from queue import Empty
from threading import Thread
from typing import Callable, Dict, List, Optional, Tuple, Union
from mmcv.utils.misc import is_method_overridden
from mmpose.utils import StopWatch
from ..utils import Message, VideoEndingMessage, limit_max_fps
@dataclass
class BufferInfo():
"""Dataclass for buffer information."""
buffer_name: str
input_name: Optional[str] = None
essential: bool = False
@dataclass
class EventInfo():
"""Dataclass for event handler information."""
event_name: str
is_keyboard: bool = False
handler_func: Optional[Callable] = None
class Node(Thread, metaclass=ABCMeta):
"""Base interface of functional module.
Parameters:
name (str, optional): The node name (also thread name).
enable_key (str|int, optional): Set a hot-key to toggle enable/disable
of the node. If an int value is given, it will be treated as an
ascii code of a key. Please note:
1. If enable_key is set, the bypass method need to be
overridden to define the node behavior when disabled
2. Some hot-key has been use for particular use. For example:
'q', 'Q' and 27 are used for quit
Default: None
max_fps (int): Maximum FPS of the node. This is to avoid the node
running unrestrictedly and causing large resource consuming.
Default: 30
input_check_interval (float): Minimum interval (in millisecond) between
checking if input is ready. Default: 0.001
enable (bool): Default enable/disable status. Default: True.
daemon (bool): Whether node is a daemon. Default: True.
"""
def __init__(self,
name: Optional[str] = None,
enable_key: Optional[Union[str, int]] = None,
max_fps: int = 30,
input_check_interval: float = 0.01,
enable: bool = True,
daemon=False):
super().__init__(name=name, daemon=daemon)
self._runner = None
self._enabled = enable
self.enable_key = enable_key
self.max_fps = max_fps
self.input_check_interval = input_check_interval
# A partitioned buffer manager the runner's buffer manager that
# only accesses the buffers related to the node
self._buffer_manager = None
# Input/output buffers are a list of registered buffers' information
self._input_buffers = []
self._output_buffers = []
# Event manager is a copy of assigned runner's event manager
self._event_manager = None
# A list of registered event information
# See register_event() for more information
# Note that we recommend to handle events in nodes by registering
# handlers, but one can still access the raw event by _event_manager
self._registered_events = []
# A list of (listener_threads, event_info)
# See set_runner() for more information
self._event_listener_threads = []
# A timer to calculate node FPS
self._timer = StopWatch(window=10)
# Register enable toggle key
if self.enable_key:
# If the node allows toggling enable, it should override the
# `bypass` method to define the node behavior when disabled.
if not is_method_overridden('bypass', Node, self.__class__):
raise NotImplementedError(
f'The node {self.__class__} does not support toggling'
'enable but got argument `enable_key`. To support toggling'
'enable, please override the `bypass` method of the node.')
self.register_event(
event_name=self.enable_key,
is_keyboard=True,
handler_func=self._toggle_enable,
)
@property
def registered_buffers(self):
return self._input_buffers + self._output_buffers
@property
def registered_events(self):
return self._registered_events.copy()
def _toggle_enable(self):
self._enabled = not self._enabled
def register_input_buffer(self,
buffer_name: str,
input_name: str,
essential: bool = False):
"""Register an input buffer, so that Node can automatically check if
data is ready, fetch data from the buffers and format the inputs to
feed into `process` method.
This method can be invoked multiple times to register multiple input
buffers.
The subclass of Node should invoke `register_input_buffer` in its
`__init__` method.
Args:
buffer_name (str): The name of the buffer
input_name (str): The name of the fetched message from the
corresponding buffer
essential (bool): An essential input means the node will wait
until the input is ready before processing. Otherwise, an
inessential input will not block the processing, instead
a None will be fetched if the buffer is not ready.
"""
buffer_info = BufferInfo(buffer_name, input_name, essential)
self._input_buffers.append(buffer_info)
def register_output_buffer(self, buffer_name: Union[str, List[str]]):
"""Register one or multiple output buffers, so that the Node can
automatically send the output of the `process` method to these buffers.
The subclass of Node should invoke `register_output_buffer` in its
`__init__` method.
Args:
buffer_name (str|list): The name(s) of the output buffer(s).
"""
if not isinstance(buffer_name, list):
buffer_name = [buffer_name]
for name in buffer_name:
buffer_info = BufferInfo(name)
self._output_buffers.append(buffer_info)
def register_event(self,
event_name: str,
is_keyboard: bool = False,
handler_func: Optional[Callable] = None):
"""Register an event. All events used in the node need to be registered
in __init__(). If a callable handler is given, a thread will be create
to listen and handle the event when the node starts.
Args:
Args:
event_name (str|int): The event name. If is_keyboard==True,
event_name should be a str (as char) or an int (as ascii)
is_keyboard (bool): Indicate whether it is an keyboard
event. If True, the argument event_name will be regarded as a
key indicator.
handler_func (callable, optional): The event handler function,
which should be a collable object with no arguments or
return values. Default: None.
"""
event_info = EventInfo(event_name, is_keyboard, handler_func)
self._registered_events.append(event_info)
def set_runner(self, runner):
# Get partitioned buffer manager
buffer_names = [
buffer.buffer_name
for buffer in self._input_buffers + self._output_buffers
]
self._buffer_manager = runner.buffer_manager.get_sub_manager(
buffer_names)
# Get event manager
self._event_manager = runner.event_manager
def _get_input_from_buffer(self) -> Tuple[bool, Optional[Dict]]:
"""Get and pack input data if it's ready. The function returns a tuple
of a status flag and a packed data dictionary. If input_buffer is
ready, the status flag will be True, and the packed data is a dict
whose items are buffer names and corresponding messages (unready
additional buffers will give a `None`). Otherwise, the status flag is
False and the packed data is None.
Returns:
bool: status flag
dict[str, Message]: the packed inputs where the key is the buffer
name and the value is the Message got from the corresponding
buffer.
"""
buffer_manager = self._buffer_manager
if buffer_manager is None:
raise ValueError(f'{self.name}: Runner not set!')
# Check that essential buffers are ready
for buffer_info in self._input_buffers:
if buffer_info.essential and buffer_manager.is_empty(
buffer_info.buffer_name):
return False, None
# Default input
result = {
buffer_info.input_name: None
for buffer_info in self._input_buffers
}
for buffer_info in self._input_buffers:
try:
result[buffer_info.input_name] = buffer_manager.get(
buffer_info.buffer_name, block=False)
except Empty:
if buffer_info.essential:
# Return unsuccessful flag if any
# essential input is unready
return False, None
return True, result
def _send_output_to_buffers(self, output_msg):
"""Send output of the process method to registered output buffers.
Args:
output_msg (Message): output message
force (bool, optional): If True, block until the output message
has been put into all output buffers. Default: False
"""
for buffer_info in self._output_buffers:
buffer_name = buffer_info.buffer_name
self._buffer_manager.put_force(buffer_name, output_msg)
@abstractmethod
def process(self, input_msgs: Dict[str, Message]) -> Union[Message, None]:
"""The core method that implement the function of the node. This method
will be invoked when the node is enabled and the input data is ready.
All subclasses of Node should override this method.
Args:
input_msgs (dict): The input data collected from the buffers. For
each item, the key is the `input_name` of the registered input
buffer, while the value is a Message instance fetched from the
buffer (or None if the buffer is unessential and not ready).
Returns:
Message: The output message of the node. It will be send to all
registered output buffers.
"""
def bypass(self, input_msgs: Dict[str, Message]) -> Union[Message, None]:
"""The method that defines the node behavior when disabled. Note that
if the node has an `enable_key`, this method should be override.
The method input/output is same as it of `process` method.
Args:
input_msgs (dict): The input data collected from the buffers. For
each item, the key is the `input_name` of the registered input
buffer, while the value is a Message instance fetched from the
buffer (or None if the buffer is unessential and not ready).
Returns:
Message: The output message of the node. It will be send to all
registered output buffers.
"""
raise NotImplementedError
def _get_node_info(self):
"""Get route information of the node."""
info = {'fps': self._timer.report('_FPS_'), 'timestamp': time.time()}
return info
def on_exit(self):
"""This method will be invoked on event `_exit_`.
Subclasses should override this method to specifying the exiting
behavior.
"""
def run(self):
"""Method representing the Node's activity.
This method override the standard run() method of Thread. Users should
not override this method in subclasses.
"""
logging.info(f'Node {self.name} starts')
# Create event listener threads
for event_info in self._registered_events:
if event_info.handler_func is None:
continue
def event_listener():
while True:
with self._event_manager.wait_and_handle(
event_info.event_name, event_info.is_keyboard):
event_info.handler_func()
t_listener = Thread(target=event_listener, args=(), daemon=True)
t_listener.start()
self._event_listener_threads.append(t_listener)
# Loop
while True:
# Exit
if self._event_manager.is_set('_exit_'):
self.on_exit()
break
# Check if input is ready
input_status, input_msgs = self._get_input_from_buffer()
# Input is not ready
if not input_status:
time.sleep(self.input_check_interval)
continue
# If a VideoEndingMessage is received, broadcast the signal
# without invoking process() or bypass()
video_ending = False
for _, msg in input_msgs.items():
if isinstance(msg, VideoEndingMessage):
self._send_output_to_buffers(msg)
video_ending = True
break
if video_ending:
self.on_exit()
break
# Check if enabled
if not self._enabled:
# Override bypass method to define node behavior when disabled
output_msg = self.bypass(input_msgs)
else:
with self._timer.timeit():
with limit_max_fps(self.max_fps):
# Process
output_msg = self.process(input_msgs)
if output_msg:
# Update route information
node_info = self._get_node_info()
output_msg.update_route_info(node=self, info=node_info)
# Send output message
if output_msg is not None:
self._send_output_to_buffers(output_msg)
logging.info(f'{self.name}: process ending.')
@@ -0,0 +1,340 @@
# Copyright (c) OpenMMLab. All rights reserved.
import time
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple, Union
import cv2
import numpy as np
from ..utils import (FrameMessage, get_eye_keypoint_ids, get_hand_keypoint_ids,
get_mouth_keypoint_ids, load_image_from_disk_or_url)
from .builder import NODES
from .frame_drawing_node import FrameDrawingNode
@dataclass
class HeartInfo():
"""Dataclass for heart information."""
heart_type: int
start_time: float
start_pos: Tuple[int, int]
end_pos: Tuple[int, int]
@NODES.register_module()
class ValentineMagicNode(FrameDrawingNode):
def __init__(self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
kpt_vis_thr: float = 0.3,
hand_heart_angle_thr: float = 90.0,
longest_duration: float = 2.0,
largest_ratio: float = 0.25,
hand_heart_img_path: Optional[str] = None,
flying_heart_img_path: Optional[str] = None,
hand_heart_dis_ratio_thr: float = 1.0,
flying_heart_dis_ratio_thr: float = 3.5,
num_persons: int = 2):
super().__init__(
name, frame_buffer, output_buffer, enable_key=enable_key)
if hand_heart_img_path is None:
hand_heart_img_path = 'https://user-images.githubusercontent.com/'\
'87690686/149731850-ea946766-a4e8-4efa-82f5'\
'-e2f0515db8ae.png'
if flying_heart_img_path is None:
flying_heart_img_path = 'https://user-images.githubusercontent.'\
'com/87690686/153554948-937ce496-33dd-4'\
'9ab-9829-0433fd7c13c4.png'
self.hand_heart = load_image_from_disk_or_url(hand_heart_img_path)
self.flying_heart = load_image_from_disk_or_url(flying_heart_img_path)
self.kpt_vis_thr = kpt_vis_thr
self.hand_heart_angle_thr = hand_heart_angle_thr
self.hand_heart_dis_ratio_thr = hand_heart_dis_ratio_thr
self.flying_heart_dis_ratio_thr = flying_heart_dis_ratio_thr
self.longest_duration = longest_duration
self.largest_ratio = largest_ratio
self.num_persons = num_persons
# record the heart infos for each person
self.heart_infos = {}
def _cal_distance(self, p1: np.ndarray, p2: np.ndarray) -> np.float64:
"""calculate the distance of points p1 and p2."""
return np.sqrt((p1[0] - p2[0])**2 + (p1[1] - p2[1])**2)
def _cal_angle(self, p1: np.ndarray, p2: np.ndarray, p3: np.ndarray,
p4: np.ndarray) -> np.float64:
"""calculate the angle of vectors v1(constructed by points p2 and p1)
and v2(constructed by points p4 and p3)"""
v1 = p2 - p1
v2 = p4 - p3
vector_prod = v1[0] * v2[0] + v1[1] * v2[1]
length_prod = np.sqrt(pow(v1[0], 2) + pow(v1[1], 2)) * np.sqrt(
pow(v2[0], 2) + pow(v2[1], 2))
cos = vector_prod * 1.0 / (length_prod * 1.0 + 1e-6)
return (np.arccos(cos) / np.pi) * 180
def _check_heart(self, pred: Dict[str,
np.ndarray], hand_indices: List[int],
mouth_index: int, eye_indices: List[int]) -> int:
"""Check the type of Valentine Magic based on the pose results and
keypoint indices of hand, mouth. and eye.
Args:
pred(dict): The pose estimation results containing:
- "keypoints" (np.ndarray[K,3]): keypoint detection result
in [x, y, score]
hand_indices(list[int]): keypoint indices of hand
mouth_index(int): keypoint index of mouth
eye_indices(list[int]): keypoint indices of eyes
Returns:
int: a number representing the type of heart pose,
0: None, 1: hand heart, 2: left hand blow kiss,
3: right hand blow kiss
"""
kpts = pred['keypoints']
left_eye_idx, right_eye_idx = eye_indices
left_eye_pos = kpts[left_eye_idx][:2]
right_eye_pos = kpts[right_eye_idx][:2]
eye_dis = self._cal_distance(left_eye_pos, right_eye_pos)
# these indices are corresoponding to the following keypoints:
# left_hand_root, left_pinky_finger1,
# left_pinky_finger3, left_pinky_finger4,
# right_hand_root, right_pinky_finger1
# right_pinky_finger3, right_pinky_finger4
both_hands_vis = True
for i in [0, 17, 19, 20, 21, 38, 40, 41]:
if kpts[hand_indices[i]][2] < self.kpt_vis_thr:
both_hands_vis = False
if both_hands_vis:
p1 = kpts[hand_indices[20]][:2]
p2 = kpts[hand_indices[19]][:2]
p3 = kpts[hand_indices[17]][:2]
p4 = kpts[hand_indices[0]][:2]
left_angle = self._cal_angle(p1, p2, p3, p4)
p1 = kpts[hand_indices[41]][:2]
p2 = kpts[hand_indices[40]][:2]
p3 = kpts[hand_indices[38]][:2]
p4 = kpts[hand_indices[21]][:2]
right_angle = self._cal_angle(p1, p2, p3, p4)
hand_dis = self._cal_distance(kpts[hand_indices[20]][:2],
kpts[hand_indices[41]][:2])
if (left_angle < self.hand_heart_angle_thr
and right_angle < self.hand_heart_angle_thr
and hand_dis / eye_dis < self.hand_heart_dis_ratio_thr):
return 1
# these indices are corresoponding to the following keypoints:
# left_middle_finger1, left_middle_finger4,
left_hand_vis = True
for i in [9, 12]:
if kpts[hand_indices[i]][2] < self.kpt_vis_thr:
left_hand_vis = False
break
# right_middle_finger1, right_middle_finger4
right_hand_vis = True
for i in [30, 33]:
if kpts[hand_indices[i]][2] < self.kpt_vis_thr:
right_hand_vis = False
break
mouth_vis = True
if kpts[mouth_index][2] < self.kpt_vis_thr:
mouth_vis = False
if (not left_hand_vis and not right_hand_vis) or not mouth_vis:
return 0
mouth_pos = kpts[mouth_index]
left_mid_hand_pos = (kpts[hand_indices[9]][:2] +
kpts[hand_indices[12]][:2]) / 2
lefthand_mouth_dis = self._cal_distance(left_mid_hand_pos, mouth_pos)
if lefthand_mouth_dis / eye_dis < self.flying_heart_dis_ratio_thr:
return 2
right_mid_hand_pos = (kpts[hand_indices[30]][:2] +
kpts[hand_indices[33]][:2]) / 2
righthand_mouth_dis = self._cal_distance(right_mid_hand_pos, mouth_pos)
if righthand_mouth_dis / eye_dis < self.flying_heart_dis_ratio_thr:
return 3
return 0
def _get_heart_route(self, heart_type: int, cur_pred: Dict[str,
np.ndarray],
tar_pred: Dict[str,
np.ndarray], hand_indices: List[int],
mouth_index: int) -> Tuple[int, int]:
"""get the start and end position of the heart, based on two keypoint
results and keypoint indices of hand and mouth.
Args:
cur_pred(dict): The pose estimation results of current person,
containing: the following keys:
- "keypoints" (np.ndarray[K,3]): keypoint detection result
in [x, y, score]
tar_pred(dict): The pose estimation results of target person,
containing: the following keys:
- "keypoints" (np.ndarray[K,3]): keypoint detection result
in [x, y, score]
hand_indices(list[int]): keypoint indices of hand
mouth_index(int): keypoint index of mouth
Returns:
tuple(int): the start position of heart
tuple(int): the end position of heart
"""
cur_kpts = cur_pred['keypoints']
assert heart_type in [1, 2,
3], 'Can not determine the type of heart effect'
if heart_type == 1:
p1 = cur_kpts[hand_indices[20]][:2]
p2 = cur_kpts[hand_indices[41]][:2]
elif heart_type == 2:
p1 = cur_kpts[hand_indices[9]][:2]
p2 = cur_kpts[hand_indices[12]][:2]
elif heart_type == 3:
p1 = cur_kpts[hand_indices[30]][:2]
p2 = cur_kpts[hand_indices[33]][:2]
cur_x, cur_y = (p1 + p2) / 2
# the mid point of two fingers
start_pos = (int(cur_x), int(cur_y))
tar_kpts = tar_pred['keypoints']
end_pos = tar_kpts[mouth_index][:2]
return start_pos, end_pos
def _draw_heart(self, canvas: np.ndarray, heart_info: HeartInfo,
t_pass: float) -> np.ndarray:
"""draw the heart according to heart info and time."""
start_x, start_y = heart_info.start_pos
end_x, end_y = heart_info.end_pos
scale = t_pass / self.longest_duration
max_h, max_w = canvas.shape[:2]
hm, wm = self.largest_ratio * max_h, self.largest_ratio * max_h
new_h, new_w = int(hm * scale), int(wm * scale)
x = int(start_x + scale * (end_x - start_x))
y = int(start_y + scale * (end_y - start_y))
y1 = max(0, y - int(new_h / 2))
y2 = min(max_h - 1, y + int(new_h / 2))
x1 = max(0, x - int(new_w / 2))
x2 = min(max_w - 1, x + int(new_w / 2))
target = canvas[y1:y2 + 1, x1:x2 + 1].copy()
new_h, new_w = target.shape[:2]
if new_h == 0 or new_w == 0:
return canvas
assert heart_info.heart_type in [
1, 2, 3
], 'Can not determine the type of heart effect'
if heart_info.heart_type == 1: # hand heart
patch = self.hand_heart.copy()
elif heart_info.heart_type >= 2: # hand blow kiss
patch = self.flying_heart.copy()
if heart_info.start_pos[0] > heart_info.end_pos[0]:
patch = patch[:, ::-1]
patch = cv2.resize(patch, (new_w, new_h))
mask = cv2.cvtColor(patch, cv2.COLOR_BGR2GRAY)
mask = (mask < 100)[..., None].astype(np.float32) * 0.8
canvas[y1:y2 + 1, x1:x2 + 1] = patch * mask + target * (1 - mask)
return canvas
def draw(self, frame_msg: FrameMessage) -> np.ndarray:
canvas = frame_msg.get_image()
pose_results = frame_msg.get_pose_results()
if not pose_results:
return canvas
for pose_result in pose_results:
model_cfg = pose_result['model_cfg']
preds = [pred.copy() for pred in pose_result['preds']]
# if number of persons in the image is less than 2,
# no heart effect will be triggered
if len(preds) < self.num_persons:
continue
# if number of persons in the image is more than 2,
# only use the first two pose results
preds = preds[:self.num_persons]
ids = [preds[i]['track_id'] for i in range(self.num_persons)]
for id in self.heart_infos.copy():
if id not in ids:
# if the id of a person not in previous heart_infos,
# delete the corresponding field
del self.heart_infos[id]
for i in range(self.num_persons):
id = preds[i]['track_id']
# if the predicted person in previous heart_infos,
# draw the heart
if id in self.heart_infos.copy():
t_pass = time.time() - self.heart_infos[id].start_time
# the time passed since last heart pose less than
# longest_duration, continue to draw the heart
if t_pass < self.longest_duration:
canvas = self._draw_heart(canvas, self.heart_infos[id],
t_pass)
# reset corresponding heart info
else:
del self.heart_infos[id]
else:
hand_indices = get_hand_keypoint_ids(model_cfg)
mouth_index = get_mouth_keypoint_ids(model_cfg)
eye_indices = get_eye_keypoint_ids(model_cfg)
# check the type of Valentine Magic based on pose results
# and keypoint indices of hand and mouth
heart_type = self._check_heart(preds[i], hand_indices,
mouth_index, eye_indices)
# trigger a Valentine Magic effect
if heart_type:
# get the route of heart
start_pos, end_pos = self._get_heart_route(
heart_type, preds[i],
preds[self.num_persons - 1 - i], hand_indices,
mouth_index)
start_time = time.time()
self.heart_infos[id] = HeartInfo(
heart_type, start_time, start_pos, end_pos)
return canvas
@@ -0,0 +1,240 @@
# Copyright (c) OpenMMLab. All rights reserved.
import json
from dataclasses import dataclass
from typing import List, Tuple, Union
import cv2
import numpy as np
from mmpose.datasets.dataset_info import DatasetInfo
from ..utils import load_image_from_disk_or_url
from .builder import NODES
from .frame_drawing_node import FrameDrawingNode
@dataclass
class DynamicInfo:
pos_curr: Tuple[int, int] = (0, 0)
pos_step: Tuple[int, int] = (0, 0)
step_curr: int = 0
@NODES.register_module()
class XDwenDwenNode(FrameDrawingNode):
"""An effect drawing node that captures the face of a cat or dog and blend
it into a Bing-Dwen-Dwen (the mascot of 2022 Beijing Winter Olympics).
Parameters:
name (str, optional): The node name (also thread name).
frame_buffer (str): The name of the input buffer.
output_buffer (str | list): The name(s) of the output buffer(s).
mode_key (str | int): A hot key to switch the background image.
resource_file (str): The annotation file of resource images, which
should be in Labelbee format and contain both facial keypoint and
region annotations.
out_shape (tuple): The shape of output frame in (width, height).
"""
dynamic_scale = 0.15
dynamic_max_step = 15
def __init__(
self,
name: str,
frame_buffer: str,
output_buffer: Union[str, List[str]],
mode_key: Union[str, int],
resource_file: str,
out_shape: Tuple[int, int] = (480, 480),
rigid_transform: bool = True,
):
super().__init__(name, frame_buffer, output_buffer, enable=True)
self.mode_key = mode_key
self.mode_index = 0
self.out_shape = out_shape
self.rigid = rigid_transform
self.latest_pred = None
self.dynamic_info = DynamicInfo()
self.register_event(
self.mode_key, is_keyboard=True, handler_func=self.switch_mode)
self._init_resource(resource_file)
def _init_resource(self, resource_file):
# The resource_file is a JSON file that contains the facial
# keypoint and mask annotation information of the resource files.
# The annotations should follow the label-bee standard format.
# See https://github.com/open-mmlab/labelbee-client for details.
with open(resource_file) as f:
anns = json.load(f)
resource_infos = []
for ann in anns:
# Load image
img = load_image_from_disk_or_url(ann['url'])
# Load result
rst = json.loads(ann['result'])
# Check facial keypoint information
assert rst['step_1']['toolName'] == 'pointTool'
assert len(rst['step_1']['result']) == 3
keypoints = sorted(
rst['step_1']['result'], key=lambda x: x['order'])
keypoints = np.array([[pt['x'], pt['y']] for pt in keypoints])
# Check facial mask
assert rst['step_2']['toolName'] == 'polygonTool'
assert len(rst['step_2']['result']) == 1
assert len(rst['step_2']['result'][0]['pointList']) > 2
mask_pts = np.array(
[[pt['x'], pt['y']]
for pt in rst['step_2']['result'][0]['pointList']])
mul = 1.0 + self.dynamic_scale
w_scale = self.out_shape[0] / img.shape[1] * mul
h_scale = self.out_shape[1] / img.shape[0] * mul
img = cv2.resize(
img,
dsize=None,
fx=w_scale,
fy=h_scale,
interpolation=cv2.INTER_CUBIC)
keypoints *= [w_scale, h_scale]
mask_pts *= [w_scale, h_scale]
mask = cv2.fillPoly(
np.zeros(img.shape[:2], dtype=np.uint8),
[mask_pts.astype(np.int32)],
color=1)
res = {
'img': img,
'keypoints': keypoints,
'mask': mask,
}
resource_infos.append(res)
self.resource_infos = resource_infos
self._reset_dynamic()
def switch_mode(self):
self.mode_index = (self.mode_index + 1) % len(self.resource_infos)
def _reset_dynamic(self):
x_tar = np.random.randint(int(self.out_shape[0] * self.dynamic_scale))
y_tar = np.random.randint(int(self.out_shape[1] * self.dynamic_scale))
x_step = (x_tar -
self.dynamic_info.pos_curr[0]) / self.dynamic_max_step
y_step = (y_tar -
self.dynamic_info.pos_curr[1]) / self.dynamic_max_step
self.dynamic_info.pos_step = (x_step, y_step)
self.dynamic_info.step_curr = 0
def draw(self, frame_msg):
full_pose_results = frame_msg.get_pose_results()
pred = None
if full_pose_results:
for pose_results in full_pose_results:
if not pose_results['preds']:
continue
pred = pose_results['preds'][0].copy()
pred['dataset'] = DatasetInfo(pose_results['model_cfg'].data.
test.dataset_info).dataset_name
self.latest_pred = pred
break
# Use the latest pose result if there is none available in
# the current frame.
if pred is None:
pred = self.latest_pred
# Get the background image and facial annotations
res = self.resource_infos[self.mode_index]
img = frame_msg.get_image()
canvas = res['img'].copy()
mask = res['mask']
kpts_tar = res['keypoints']
if pred is not None:
if pred['dataset'] == 'ap10k':
# left eye: 0, right eye: 1, nose: 2
kpts_src = pred['keypoints'][[0, 1, 2], :2]
elif pred['dataset'] == 'coco_wholebody':
# left eye: 1, right eye 2, nose: 0
kpts_src = pred['keypoints'][[1, 2, 0], :2]
else:
raise ValueError('Can not obtain face landmark information'
f'from dataset: {pred["type"]}')
trans_mat = self._get_transform(kpts_src, kpts_tar)
warp = cv2.warpAffine(img, trans_mat, dsize=canvas.shape[:2])
cv2.copyTo(warp, mask, canvas)
# Add random movement to the background
xc, yc = self.dynamic_info.pos_curr
xs, ys = self.dynamic_info.pos_step
w, h = self.out_shape
x = min(max(int(xc), 0), canvas.shape[1] - w + 1)
y = min(max(int(yc), 0), canvas.shape[0] - h + 1)
canvas = canvas[y:y + h, x:x + w]
self.dynamic_info.pos_curr = (xc + xs, yc + ys)
self.dynamic_info.step_curr += 1
if self.dynamic_info.step_curr == self.dynamic_max_step:
self._reset_dynamic()
return canvas
def _get_transform(self, kpts_src, kpts_tar):
if self.rigid:
# rigid transform
n = kpts_src.shape[0]
X = np.zeros((n * 2, 4), dtype=np.float32)
U = np.zeros((n * 2, 1), dtype=np.float32)
X[:n, :2] = kpts_src
X[:n, 2] = 1
X[n:, 0] = kpts_src[:, 1]
X[n:, 1] = -kpts_src[:, 0]
X[n:, 3] = 1
U[:n, 0] = kpts_tar[:, 0]
U[n:, 0] = kpts_tar[:, 1]
M = np.linalg.pinv(X).dot(U).flatten()
trans_mat = np.array([[M[0], M[1], M[2]], [-M[1], M[0], M[3]]],
dtype=np.float32)
else:
# normal affine transform
# adaptive horizontal flipping
if (np.linalg.norm(kpts_tar[0] - kpts_tar[2]) -
np.linalg.norm(kpts_tar[1] - kpts_tar[2])) * (
np.linalg.norm(kpts_src[0] - kpts_src[2]) -
np.linalg.norm(kpts_src[1] - kpts_src[2])) < 0:
kpts_src = kpts_src[[1, 0, 2], :]
trans_mat, _ = cv2.estimateAffine2D(
kpts_src.astype(np.float32), kpts_tar.astype(np.float32))
return trans_mat
@@ -0,0 +1,31 @@
# Copyright (c) OpenMMLab. All rights reserved.
from .buffer import BufferManager
from .event import EventManager
from .message import FrameMessage, Message, VideoEndingMessage
from .misc import (ImageCapture, copy_and_paste, expand_and_clamp,
get_cached_file_path, is_image_file, limit_max_fps,
load_image_from_disk_or_url, screen_matting)
from .pose import (get_eye_keypoint_ids, get_face_keypoint_ids,
get_hand_keypoint_ids, get_mouth_keypoint_ids,
get_wrist_keypoint_ids)
__all__ = [
'BufferManager',
'EventManager',
'FrameMessage',
'Message',
'limit_max_fps',
'VideoEndingMessage',
'load_image_from_disk_or_url',
'get_cached_file_path',
'screen_matting',
'expand_and_clamp',
'copy_and_paste',
'is_image_file',
'ImageCapture',
'get_eye_keypoint_ids',
'get_face_keypoint_ids',
'get_wrist_keypoint_ids',
'get_mouth_keypoint_ids',
'get_hand_keypoint_ids',
]
@@ -0,0 +1,106 @@
# Copyright (c) OpenMMLab. All rights reserved.
from functools import wraps
from queue import Queue
from typing import Dict, List, Optional
from mmcv import is_seq_of
__all__ = ['BufferManager']
def check_buffer_registered(exist=True):
def wrapper(func):
@wraps(func)
def wrapped(manager, name, *args, **kwargs):
if exist:
# Assert buffer exist
if name not in manager:
raise ValueError(f'Fail to call {func.__name__}: '
f'buffer "{name}" is not registered.')
else:
# Assert buffer not exist
if name in manager:
raise ValueError(f'Fail to call {func.__name__}: '
f'buffer "{name}" is already registered.')
return func(manager, name, *args, **kwargs)
return wrapped
return wrapper
class Buffer(Queue):
def put_force(self, item):
"""Force to put an item into the buffer.
If the buffer is already full, the earliest item in the buffer will be
remove to make room for the incoming item.
"""
with self.mutex:
if self.maxsize > 0:
while self._qsize() >= self.maxsize:
_ = self._get()
self.unfinished_tasks -= 1
self._put(item)
self.unfinished_tasks += 1
self.not_empty.notify()
class BufferManager():
def __init__(self,
buffer_type: type = Buffer,
buffers: Optional[Dict] = None):
self.buffer_type = buffer_type
if buffers is None:
self._buffers = {}
else:
if is_seq_of(list(buffers.values()), buffer_type):
self._buffers = buffers.copy()
else:
raise ValueError('The values of buffers should be instance '
f'of {buffer_type}')
def __contains__(self, name):
return name in self._buffers
@check_buffer_registered(False)
def register_buffer(self, name, maxsize=0):
self._buffers[name] = self.buffer_type(maxsize)
@check_buffer_registered()
def put(self, name, item, block=True, timeout=None):
self._buffers[name].put(item, block, timeout)
@check_buffer_registered()
def put_force(self, name, item):
self._buffers[name].put_force(item)
@check_buffer_registered()
def get(self, name, block=True, timeout=None):
return self._buffers[name].get(block, timeout)
@check_buffer_registered()
def is_empty(self, name):
return self._buffers[name].empty()
@check_buffer_registered()
def is_full(self, name):
return self._buffers[name].full()
def get_sub_manager(self, buffer_names: List[str]):
buffers = {name: self._buffers[name] for name in buffer_names}
return BufferManager(self.buffer_type, buffers)
def get_info(self):
buffer_info = {}
for name, buffer in self._buffers.items():
buffer_info[name] = {
'size': buffer.size,
'maxsize': buffer.maxsize
}
return buffer_info
@@ -0,0 +1,59 @@
# Copyright (c) OpenMMLab. All rights reserved.
from collections import defaultdict
from contextlib import contextmanager
from threading import Event
from typing import Optional
class EventManager():
def __init__(self):
self._events = defaultdict(Event)
def register_event(self,
event_name: str = None,
is_keyboard: bool = False):
if is_keyboard:
event_name = self._get_keyboard_event_name(event_name)
self._events[event_name] = Event()
def set(self, event_name: str = None, is_keyboard: bool = False):
if is_keyboard:
event_name = self._get_keyboard_event_name(event_name)
return self._events[event_name].set()
def wait(self,
event_name: str = None,
is_keyboard: Optional[bool] = False,
timeout: Optional[float] = None):
if is_keyboard:
event_name = self._get_keyboard_event_name(event_name)
return self._events[event_name].wait(timeout)
def is_set(self,
event_name: str = None,
is_keyboard: Optional[bool] = False):
if is_keyboard:
event_name = self._get_keyboard_event_name(event_name)
return self._events[event_name].is_set()
def clear(self,
event_name: str = None,
is_keyboard: Optional[bool] = False):
if is_keyboard:
event_name = self._get_keyboard_event_name(event_name)
return self._events[event_name].clear()
@staticmethod
def _get_keyboard_event_name(key):
return f'_keyboard_{chr(key) if isinstance(key,int) else key}'
@contextmanager
def wait_and_handle(self,
event_name: str = None,
is_keyboard: Optional[bool] = False):
self.wait(event_name, is_keyboard)
try:
yield
finally:
self.clear(event_name, is_keyboard)
@@ -0,0 +1,204 @@
# Copyright (c) OpenMMLab. All rights reserved.
import time
import uuid
import warnings
from typing import Dict, List, Optional
import numpy as np
class Message():
"""Message base class.
All message class should inherit this class. The basic use of a Message
instance is to carray a piece of text message (self.msg) and a dict that
stores structured data (self.data), e.g. frame image, model prediction,
et al.
A message may also hold route information, which is composed of
information of all nodes the message has passed through.
Parameters:
msg (str): The text message.
data (dict, optional): The structured data.
"""
def __init__(self, msg: str = '', data: Optional[Dict] = None):
self.msg = msg
self.data = data if data else {}
self.route_info = []
self.timestamp = time.time()
self.id = uuid.uuid4()
def update_route_info(self,
node=None,
node_name: Optional[str] = None,
node_type: Optional[str] = None,
info: Optional[Dict] = None):
"""Append new node information to the route information.
Args:
node (Node, optional): An instance of Node that provides basic
information like the node name and type. Default: None.
node_name (str, optional): The node name. If node is given,
node_name will be ignored. Default: None.
node_type (str, optional): The class name of the node. If node
is given, node_type will be ignored. Default: None.
info (dict, optional): The node information, which is usually
given by node.get_node_info(). Default: None.
"""
if node is not None:
if node_name is not None or node_type is not None:
warnings.warn(
'`node_name` and `node_type` will be overridden if node'
'is provided.')
node_name = node.name
node_type = node.__class__.__name__
node_info = {'node': node_name, 'node_type': node_type, 'info': info}
self.route_info.append(node_info)
def set_route_info(self, route_info: List):
"""Directly set the entire route information.
Args:
route_info (list): route information to set to the message.
"""
self.route_info = route_info
def merge_route_info(self, route_info: List):
"""Merge the given route information into the original one of the
message. This is used for combining route information from multiple
messages. The node information in the route will be reordered according
to their timestamps.
Args:
route_info (list): route information to merge.
"""
self.route_info += route_info
self.route_info.sort(key=lambda x: x.get('timestamp', np.inf))
def get_route_info(self) -> List:
return self.route_info.copy()
class VideoEndingMessage(Message):
"""A special message to indicate the input video is ending."""
class FrameMessage(Message):
"""The message to store information of a video frame.
A FrameMessage instance usually holds following data in self.data:
- image (array): The frame image
- detection_results (list): A list to hold detection results of
multiple detectors. Each element is a tuple (tag, result)
- pose_results (list): A list to hold pose estimation results of
multiple pose estimator. Each element is a tuple (tag, result)
"""
def __init__(self, img):
super().__init__(data=dict(image=img))
def get_image(self):
"""Get the frame image.
Returns:
array: The frame image.
"""
return self.data.get('image', None)
def set_image(self, img):
"""Set the frame image to the message."""
self.data['image'] = img
def add_detection_result(self, result, tag: str = None):
"""Add the detection result from one model into the message's
detection_results.
Args:
tag (str, optional): Give a tag to the result, which can be used
to retrieve specific results.
"""
if 'detection_results' not in self.data:
self.data['detection_results'] = []
self.data['detection_results'].append((tag, result))
def get_detection_results(self, tag: str = None):
"""Get detection results of the message.
Args:
tag (str, optional): If given, only the results with the tag
will be retrieved. Otherwise all results will be retrieved.
Default: None.
Returns:
list[dict]: The retrieved detection results
"""
if 'detection_results' not in self.data:
return None
if tag is None:
results = [res for _, res in self.data['detection_results']]
else:
results = [
res for _tag, res in self.data['detection_results']
if _tag == tag
]
return results
def add_pose_result(self, result, tag=None):
"""Add the pose estimation result from one model into the message's
pose_results.
Args:
tag (str, optional): Give a tag to the result, which can be used
to retrieve specific results.
"""
if 'pose_results' not in self.data:
self.data['pose_results'] = []
self.data['pose_results'].append((tag, result))
def get_pose_results(self, tag=None):
"""Get pose estimation results of the message.
Args:
tag (str, optional): If given, only the results with the tag
will be retrieved. Otherwise all results will be retrieved.
Default: None.
Returns:
list[dict]: The retrieved pose results
"""
if 'pose_results' not in self.data:
return None
if tag is None:
results = [res for _, res in self.data['pose_results']]
else:
results = [
res for _tag, res in self.data['pose_results'] if _tag == tag
]
return results
def get_full_results(self):
"""Get all model predictions of the message.
See set_full_results() for inference.
Returns:
dict: All model predictions, including:
- detection_results
- pose_results
"""
result_keys = ['detection_results', 'pose_results']
results = {k: self.data[k] for k in result_keys}
return results
def set_full_results(self, results):
"""Set full model results directly.
Args:
results (dict): All model predictions including:
- detection_results (list): see also add_detection_results()
- pose_results (list): see also add_pose_results()
"""
self.data.update(results)
@@ -0,0 +1,343 @@
# Copyright (c) OpenMMLab. All rights reserved.
import os
import os.path as osp
import sys
import time
from contextlib import contextmanager
from typing import Optional
from urllib.parse import urlparse
from urllib.request import urlopen
import cv2
import numpy as np
from torch.hub import HASH_REGEX, download_url_to_file
@contextmanager
def limit_max_fps(fps: Optional[float]):
t_start = time.time()
try:
yield
finally:
t_end = time.time()
if fps is not None:
t_sleep = 1.0 / fps - t_end + t_start
if t_sleep > 0:
time.sleep(t_sleep)
def _is_url(filename):
"""Check if the file is a url link.
Args:
filename (str): the file name or url link.
Returns:
bool: is url or not.
"""
prefixes = ['http://', 'https://']
for p in prefixes:
if filename.startswith(p):
return True
return False
def load_image_from_disk_or_url(filename, readFlag=cv2.IMREAD_COLOR):
"""Load an image file, from disk or url.
Args:
filename (str): file name on the disk or url link.
readFlag (int): readFlag for imdecode.
Returns:
np.ndarray: A loaded image
"""
if _is_url(filename):
# download the image, convert it to a NumPy array, and then read
# it into OpenCV format
resp = urlopen(filename)
image = np.asarray(bytearray(resp.read()), dtype='uint8')
image = cv2.imdecode(image, readFlag)
return image
else:
image = cv2.imread(filename, readFlag)
return image
def mkdir_or_exist(dir_name, mode=0o777):
if dir_name == '':
return
dir_name = osp.expanduser(dir_name)
os.makedirs(dir_name, mode=mode, exist_ok=True)
def get_cached_file_path(url,
save_dir=None,
progress=True,
check_hash=False,
file_name=None):
r"""Loads the Torch serialized object at the given URL.
If downloaded file is a zip file, it will be automatically decompressed
If the object is already present in `model_dir`, it's deserialized and
returned.
The default value of ``model_dir`` is ``<hub_dir>/checkpoints`` where
``hub_dir`` is the directory returned by :func:`~torch.hub.get_dir`.
Args:
url (str): URL of the object to download
save_dir (str, optional): directory in which to save the object
progress (bool, optional): whether or not to display a progress bar
to stderr. Default: True
check_hash(bool, optional): If True, the filename part of the URL
should follow the naming convention ``filename-<sha256>.ext``
where ``<sha256>`` is the first eight or more digits of the
SHA256 hash of the contents of the file. The hash is used to
ensure unique names and to verify the contents of the file.
Default: False
file_name (str, optional): name for the downloaded file. Filename
from ``url`` will be used if not set. Default: None.
"""
if save_dir is None:
save_dir = os.path.join('webcam_resources')
mkdir_or_exist(save_dir)
parts = urlparse(url)
filename = os.path.basename(parts.path)
if file_name is not None:
filename = file_name
cached_file = os.path.join(save_dir, filename)
if not os.path.exists(cached_file):
sys.stderr.write('Downloading: "{}" to {}\n'.format(url, cached_file))
hash_prefix = None
if check_hash:
r = HASH_REGEX.search(filename) # r is Optional[Match[str]]
hash_prefix = r.group(1) if r else None
download_url_to_file(url, cached_file, hash_prefix, progress=progress)
return cached_file
def screen_matting(img, color_low=None, color_high=None, color=None):
"""Screen Matting.
Args:
img (np.ndarray): Image data.
color_low (tuple): Lower limit (b, g, r).
color_high (tuple): Higher limit (b, g, r).
color (str): Support colors include:
- 'green' or 'g'
- 'blue' or 'b'
- 'black' or 'k'
- 'white' or 'w'
"""
if color_high is None or color_low is None:
if color is not None:
if color.lower() == 'g' or color.lower() == 'green':
color_low = (0, 200, 0)
color_high = (60, 255, 60)
elif color.lower() == 'b' or color.lower() == 'blue':
color_low = (230, 0, 0)
color_high = (255, 40, 40)
elif color.lower() == 'k' or color.lower() == 'black':
color_low = (0, 0, 0)
color_high = (40, 40, 40)
elif color.lower() == 'w' or color.lower() == 'white':
color_low = (230, 230, 230)
color_high = (255, 255, 255)
else:
NotImplementedError(f'Not supported color: {color}.')
else:
ValueError('color or color_high | color_low should be given.')
mask = cv2.inRange(img, np.array(color_low), np.array(color_high)) == 0
return mask.astype(np.uint8)
def expand_and_clamp(box, im_shape, s=1.25):
"""Expand the bbox and clip it to fit the image shape.
Args:
box (list): x1, y1, x2, y2
im_shape (ndarray): image shape (h, w, c)
s (float): expand ratio
Returns:
list: x1, y1, x2, y2
"""
x1, y1, x2, y2 = box[:4]
w = x2 - x1
h = y2 - y1
deta_w = w * (s - 1) / 2
deta_h = h * (s - 1) / 2
x1, y1, x2, y2 = x1 - deta_w, y1 - deta_h, x2 + deta_w, y2 + deta_h
img_h, img_w = im_shape[:2]
x1 = min(max(0, int(x1)), img_w - 1)
y1 = min(max(0, int(y1)), img_h - 1)
x2 = min(max(0, int(x2)), img_w - 1)
y2 = min(max(0, int(y2)), img_h - 1)
return [x1, y1, x2, y2]
def _find_connected_components(mask):
"""Find connected components and sort with areas.
Args:
mask (ndarray): instance segmentation result.
Returns:
ndarray (N, 5): Each item contains (x, y, w, h, area).
"""
num, labels, stats, centroids = cv2.connectedComponentsWithStats(mask)
stats = stats[stats[:, 4].argsort()]
return stats
def _find_bbox(mask):
"""Find the bounding box for the mask.
Args:
mask (ndarray): Mask.
Returns:
list(4, ): Returned box (x1, y1, x2, y2).
"""
mask_shape = mask.shape
if len(mask_shape) == 3:
assert mask_shape[-1] == 1, 'the channel of the mask should be 1.'
elif len(mask_shape) == 2:
pass
else:
NotImplementedError()
h, w = mask_shape[:2]
mask_w = mask.sum(0)
mask_h = mask.sum(1)
left = 0
right = w - 1
up = 0
down = h - 1
for i in range(w):
if mask_w[i] > 0:
break
left += 1
for i in range(w - 1, left, -1):
if mask_w[i] > 0:
break
right -= 1
for i in range(h):
if mask_h[i] > 0:
break
up += 1
for i in range(h - 1, up, -1):
if mask_h[i] > 0:
break
down -= 1
return [left, up, right, down]
def copy_and_paste(img,
background_img,
mask,
bbox=None,
effect_region=(0.2, 0.2, 0.8, 0.8),
min_size=(20, 20)):
"""Copy the image region and paste to the background.
Args:
img (np.ndarray): Image data.
background_img (np.ndarray): Background image data.
mask (ndarray): instance segmentation result.
bbox (ndarray): instance bbox, (x1, y1, x2, y2).
effect_region (tuple(4, )): The region to apply mask, the coordinates
are normalized (x1, y1, x2, y2).
"""
background_img = background_img.copy()
background_h, background_w = background_img.shape[:2]
region_h = (effect_region[3] - effect_region[1]) * background_h
region_w = (effect_region[2] - effect_region[0]) * background_w
region_aspect_ratio = region_w / region_h
if bbox is None:
bbox = _find_bbox(mask)
instance_w = bbox[2] - bbox[0]
instance_h = bbox[3] - bbox[1]
if instance_w > min_size[0] and instance_h > min_size[1]:
aspect_ratio = instance_w / instance_h
if region_aspect_ratio > aspect_ratio:
resize_rate = region_h / instance_h
else:
resize_rate = region_w / instance_w
mask_inst = mask[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])]
img_inst = img[int(bbox[1]):int(bbox[3]), int(bbox[0]):int(bbox[2])]
img_inst = cv2.resize(img_inst, (int(
resize_rate * instance_w), int(resize_rate * instance_h)))
mask_inst = cv2.resize(
mask_inst,
(int(resize_rate * instance_w), int(resize_rate * instance_h)),
interpolation=cv2.INTER_NEAREST)
mask_ids = list(np.where(mask_inst == 1))
mask_ids[1] += int(effect_region[0] * background_w)
mask_ids[0] += int(effect_region[1] * background_h)
background_img[tuple(mask_ids)] = img_inst[np.where(mask_inst == 1)]
return background_img
def is_image_file(path):
if isinstance(path, str):
if path.lower().endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp')):
return True
return False
class ImageCapture:
"""A mock-up version of cv2.VideoCapture that always return a const image.
Args:
image (str | ndarray): The image or image path
"""
def __init__(self, image):
if isinstance(image, str):
self.image = load_image_from_disk_or_url(image)
else:
self.image = image
def isOpened(self):
return (self.image is not None)
def read(self):
return True, self.image.copy()
def release(self):
pass
def get(self, propId):
if propId == cv2.CAP_PROP_FRAME_WIDTH:
return self.image.shape[1]
elif propId == cv2.CAP_PROP_FRAME_HEIGHT:
return self.image.shape[0]
elif propId == cv2.CAP_PROP_FPS:
return np.nan
else:
raise NotImplementedError()
@@ -0,0 +1,226 @@
# Copyright (c) OpenMMLab. All rights reserved.
from typing import List, Tuple
from mmcv import Config
from mmpose.datasets.dataset_info import DatasetInfo
def get_eye_keypoint_ids(model_cfg: Config) -> Tuple[int, int]:
"""A helpfer function to get the keypoint indices of left and right eyes
from the model config.
Args:
model_cfg (Config): pose model config.
Returns:
int: left eye keypoint index.
int: right eye keypoint index.
"""
left_eye_idx = None
right_eye_idx = None
# try obtaining eye point ids from dataset_info
try:
dataset_info = DatasetInfo(model_cfg.data.test.dataset_info)
left_eye_idx = dataset_info.keypoint_name2id.get('left_eye', None)
right_eye_idx = dataset_info.keypoint_name2id.get('right_eye', None)
except AttributeError:
left_eye_idx = None
right_eye_idx = None
if left_eye_idx is None or right_eye_idx is None:
# Fall back to hard coded keypoint id
dataset_name = model_cfg.data.test.type
if dataset_name in {
'TopDownCocoDataset', 'TopDownCocoWholeBodyDataset'
}:
left_eye_idx = 1
right_eye_idx = 2
elif dataset_name in {'AnimalPoseDataset', 'AnimalAP10KDataset'}:
left_eye_idx = 0
right_eye_idx = 1
else:
raise ValueError('Can not determine the eye keypoint id of '
f'{dataset_name}')
return left_eye_idx, right_eye_idx
def get_face_keypoint_ids(model_cfg: Config) -> Tuple[int, int]:
"""A helpfer function to get the keypoint indices of the face from the
model config.
Args:
model_cfg (Config): pose model config.
Returns:
list[int]: face keypoint index.
"""
face_indices = None
# try obtaining nose point ids from dataset_info
try:
dataset_info = DatasetInfo(model_cfg.data.test.dataset_info)
for id in range(68):
face_indices.append(
dataset_info.keypoint_name2id.get(f'face_{id}', None))
except AttributeError:
face_indices = None
if face_indices is None:
# Fall back to hard coded keypoint id
dataset_name = model_cfg.data.test.type
if dataset_name in {'TopDownCocoWholeBodyDataset'}:
face_indices = list(range(23, 91))
else:
raise ValueError('Can not determine the face id of '
f'{dataset_name}')
return face_indices
def get_wrist_keypoint_ids(model_cfg: Config) -> Tuple[int, int]:
"""A helpfer function to get the keypoint indices of left and right wrist
from the model config.
Args:
model_cfg (Config): pose model config.
Returns:
int: left wrist keypoint index.
int: right wrist keypoint index.
"""
# try obtaining eye point ids from dataset_info
try:
dataset_info = DatasetInfo(model_cfg.data.test.dataset_info)
left_wrist_idx = dataset_info.keypoint_name2id.get('left_wrist', None)
right_wrist_idx = dataset_info.keypoint_name2id.get(
'right_wrist', None)
except AttributeError:
left_wrist_idx = None
right_wrist_idx = None
if left_wrist_idx is None or right_wrist_idx is None:
# Fall back to hard coded keypoint id
dataset_name = model_cfg.data.test.type
if dataset_name in {
'TopDownCocoDataset', 'TopDownCocoWholeBodyDataset'
}:
left_wrist_idx = 9
right_wrist_idx = 10
elif dataset_name == 'AnimalPoseDataset':
left_wrist_idx = 16
right_wrist_idx = 17
elif dataset_name == 'AnimalAP10KDataset':
left_wrist_idx = 7
right_wrist_idx = 10
else:
raise ValueError('Can not determine the eye keypoint id of '
f'{dataset_name}')
return left_wrist_idx, right_wrist_idx
def get_mouth_keypoint_ids(model_cfg: Config) -> Tuple[int, int]:
"""A helpfer function to get the keypoint indices of the left and right
part of mouth from the model config.
Args:
model_cfg (Config): pose model config.
Returns:
int: left-part mouth keypoint index.
int: right-part mouth keypoint index.
"""
# try obtaining mouth point ids from dataset_info
try:
dataset_info = DatasetInfo(model_cfg.data.test.dataset_info)
mouth_index = dataset_info.keypoint_name2id.get('face-62', None)
except AttributeError:
mouth_index = None
if mouth_index is None:
# Fall back to hard coded keypoint id
dataset_name = model_cfg.data.test.type
if dataset_name == 'TopDownCocoWholeBodyDataset':
mouth_index = 85
else:
raise ValueError('Can not determine the eye keypoint id of '
f'{dataset_name}')
return mouth_index
def get_hand_keypoint_ids(model_cfg: Config) -> List[int]:
"""A helpfer function to get the keypoint indices of left and right hand
from the model config.
Args:
model_cfg (Config): pose model config.
Returns:
list[int]: hand keypoint indices.
"""
# try obtaining hand keypoint ids from dataset_info
try:
hand_indices = []
dataset_info = DatasetInfo(model_cfg.data.test.dataset_info)
hand_indices.append(
dataset_info.keypoint_name2id.get('left_hand_root', None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'left_thumb{id}', None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'left_forefinger{id}',
None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'left_middle_finger{id}',
None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'left_ring_finger{id}',
None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'left_pinky_finger{id}',
None))
hand_indices.append(
dataset_info.keypoint_name2id.get('right_hand_root', None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'right_thumb{id}', None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'right_forefinger{id}',
None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'right_middle_finger{id}',
None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'right_ring_finger{id}',
None))
for id in range(1, 5):
hand_indices.append(
dataset_info.keypoint_name2id.get(f'right_pinky_finger{id}',
None))
except AttributeError:
hand_indices = None
if hand_indices is None:
# Fall back to hard coded keypoint id
dataset_name = model_cfg.data.test.type
if dataset_name in {'TopDownCocoWholeBodyDataset'}:
hand_indices = list(range(91, 133))
else:
raise ValueError('Can not determine the hand id of '
f'{dataset_name}')
return hand_indices
@@ -0,0 +1,272 @@
# Copyright (c) OpenMMLab. All rights reserved.
import logging
import sys
import time
import warnings
from contextlib import nullcontext
from threading import Thread
from typing import Dict, List, Optional, Tuple, Union
import cv2
from .nodes import NODES
from .utils import (BufferManager, EventManager, FrameMessage, ImageCapture,
VideoEndingMessage, is_image_file, limit_max_fps)
DEFAULT_FRAME_BUFFER_SIZE = 1
DEFAULT_INPUT_BUFFER_SIZE = 1
DEFAULT_DISPLAY_BUFFER_SIZE = 0
DEFAULT_USER_BUFFER_SIZE = 1
class WebcamRunner():
"""An interface for building webcam application from config.
Parameters:
name (str): Runner name.
camera_id (int | str): The camera ID (usually the ID of the default
camera is 0). Alternatively a file path or a URL can be given
to load from a video or image file.
camera_frame_shape (tuple, optional): Set the frame shape of the
camera in (width, height). If not given, the default frame shape
will be used. This argument is only valid when using a camera
as the input source. Default: None
camera_fps (int): Video reading maximum FPS. Default: 30
buffer_sizes (dict, optional): A dict to specify buffer sizes. The
key is the buffer name and the value is the buffer size.
Default: None
nodes (list): Node configs.
"""
def __init__(self,
name: str = 'Default Webcam Runner',
camera_id: Union[int, str] = 0,
camera_fps: int = 30,
camera_frame_shape: Optional[Tuple[int, int]] = None,
synchronous: bool = False,
buffer_sizes: Optional[Dict[str, int]] = None,
nodes: Optional[List[Dict]] = None):
# Basic parameters
self.name = name
self.camera_id = camera_id
self.camera_fps = camera_fps
self.camera_frame_shape = camera_frame_shape
self.synchronous = synchronous
# self.buffer_manager manages data flow between runner and nodes
self.buffer_manager = BufferManager()
# self.event_manager manages event-based asynchronous communication
self.event_manager = EventManager()
# self.node_list holds all node instance
self.node_list = []
# self.vcap is used to read camera frames. It will be built when the
# runner starts running
self.vcap = None
# Register runner events
self.event_manager.register_event('_exit_', is_keyboard=False)
if self.synchronous:
self.event_manager.register_event('_idle_', is_keyboard=False)
# Register nodes
if not nodes:
raise ValueError('No node is registered to the runner.')
# Register default buffers
if buffer_sizes is None:
buffer_sizes = {}
# _frame_ buffer
frame_buffer_size = buffer_sizes.get('_frame_',
DEFAULT_FRAME_BUFFER_SIZE)
self.buffer_manager.register_buffer('_frame_', frame_buffer_size)
# _input_ buffer
input_buffer_size = buffer_sizes.get('_input_',
DEFAULT_INPUT_BUFFER_SIZE)
self.buffer_manager.register_buffer('_input_', input_buffer_size)
# _display_ buffer
display_buffer_size = buffer_sizes.get('_display_',
DEFAULT_DISPLAY_BUFFER_SIZE)
self.buffer_manager.register_buffer('_display_', display_buffer_size)
# Build all nodes:
for node_cfg in nodes:
logging.info(f'Create node: {node_cfg.name}({node_cfg.type})')
node = NODES.build(node_cfg)
# Register node
self.node_list.append(node)
# Register buffers
for buffer_info in node.registered_buffers:
buffer_name = buffer_info.buffer_name
if buffer_name in self.buffer_manager:
continue
buffer_size = buffer_sizes.get(buffer_name,
DEFAULT_USER_BUFFER_SIZE)
self.buffer_manager.register_buffer(buffer_name, buffer_size)
logging.info(
f'Register user buffer: {buffer_name}({buffer_size})')
# Register events
for event_info in node.registered_events:
self.event_manager.register_event(
event_name=event_info.event_name,
is_keyboard=event_info.is_keyboard)
logging.info(f'Register event: {event_info.event_name}')
# Set runner for nodes
# This step is performed after node building when the runner has
# create full buffer/event managers and can
for node in self.node_list:
logging.info(f'Set runner for node: {node.name})')
node.set_runner(self)
def _read_camera(self):
"""Continually read video frames and put them into buffers."""
camera_id = self.camera_id
fps = self.camera_fps
# Build video capture
if is_image_file(camera_id):
self.vcap = ImageCapture(camera_id)
else:
self.vcap = cv2.VideoCapture(camera_id)
if self.camera_frame_shape is not None:
width, height = self.camera_frame_shape
self.vcap.set(cv2.CAP_PROP_FRAME_WIDTH, width)
self.vcap.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
if not self.vcap.isOpened():
warnings.warn(f'Cannot open camera (ID={camera_id})')
sys.exit()
# Read video frames in a loop
first_frame = True
while not self.event_manager.is_set('_exit_'):
if self.synchronous:
if first_frame:
cm = nullcontext()
else:
# Read a new frame until the last frame has been processed
cm = self.event_manager.wait_and_handle('_idle_')
else:
# Read frames with a maximum FPS
cm = limit_max_fps(fps)
first_frame = False
with cm:
# Read a frame
ret_val, frame = self.vcap.read()
if ret_val:
# Put frame message (for display) into buffer `_frame_`
frame_msg = FrameMessage(frame)
self.buffer_manager.put('_frame_', frame_msg)
# Put input message (for model inference or other use)
# into buffer `_input_`
input_msg = FrameMessage(frame.copy())
input_msg.update_route_info(
node_name='Camera Info',
node_type='dummy',
info=self._get_camera_info())
self.buffer_manager.put_force('_input_', input_msg)
else:
# Put a video ending signal
self.buffer_manager.put('_frame_', VideoEndingMessage())
self.vcap.release()
def _display(self):
"""Continually obtain and display output frames."""
output_msg = None
while not self.event_manager.is_set('_exit_'):
while self.buffer_manager.is_empty('_display_'):
time.sleep(0.001)
# Set _idle_ to allow reading next frame
if self.synchronous:
self.event_manager.set('_idle_')
# acquire output from buffer
output_msg = self.buffer_manager.get('_display_')
# None indicates input stream ends
if isinstance(output_msg, VideoEndingMessage):
self.event_manager.set('_exit_')
break
img = output_msg.get_image()
# show in a window
cv2.imshow(self.name, img)
# handle keyboard input
key = cv2.waitKey(1)
if key != -1:
self._on_keyboard_input(key)
cv2.destroyAllWindows()
def _on_keyboard_input(self, key):
"""Handle the keyboard input."""
if key in (27, ord('q'), ord('Q')):
logging.info(f'Exit event captured: {key}')
self.event_manager.set('_exit_')
else:
logging.info(f'Keyboard event captured: {key}')
self.event_manager.set(key, is_keyboard=True)
def _get_camera_info(self):
"""Return the camera information in a dict."""
frame_width = self.vcap.get(cv2.CAP_PROP_FRAME_WIDTH)
frame_height = self.vcap.get(cv2.CAP_PROP_FRAME_HEIGHT)
frame_rate = self.vcap.get(cv2.CAP_PROP_FPS)
cam_info = {
'Camera ID': self.camera_id,
'Source resolution': f'{frame_width}x{frame_height}',
'Source FPS': frame_rate,
}
return cam_info
def run(self):
"""Program entry.
This method starts all nodes as well as video I/O in separate threads.
"""
try:
# Start node threads
non_daemon_nodes = []
for node in self.node_list:
node.start()
if not node.daemon:
non_daemon_nodes.append(node)
# Create a thread to read video frames
t_read = Thread(target=self._read_camera, args=())
t_read.start()
# Run display in the main thread
self._display()
logging.info('Display shut down')
# joint non-daemon nodes and runner threads
logging.info('Camera reading about to join')
t_read.join()
for node in non_daemon_nodes:
logging.info(f'Node {node.name} about to join')
node.join()
except KeyboardInterrupt:
pass