Spaces:
Runtime error
Runtime error
File size: 4,636 Bytes
cc0dd3c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
# Copyright (c) OpenMMLab. All rights reserved.
from itertools import groupby
from typing import Dict, List, Optional, Union
import cv2
import numpy as np
from ...utils import get_eye_keypoint_ids
from ..base_visualizer_node import BaseVisualizerNode
from ..registry import NODES
@NODES.register_module()
class BigeyeEffectNode(BaseVisualizerNode):
"""Apply big-eye effect to the objects with eye keypoints in the frame.
Args:
name (str): The node name (also thread name)
input_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-keys are reserved for
particular use. For example: 'q', 'Q' and 27 are used for exiting.
Default: ``None``
enable (bool): Default enable/disable status. Default: ``True``
kpt_thr (float): The score threshold of valid keypoints. Default: 0.5
Example::
>>> cfg = dict(
... type='SunglassesEffectNode',
... name='sunglasses',
... enable_key='s',
... enable=False,
... input_buffer='vis',
... output_buffer='vis_sunglasses')
>>> from mmpose.apis.webcam.nodes import NODES
>>> node = NODES.build(cfg)
"""
def __init__(self,
name: str,
input_buffer: str,
output_buffer: Union[str, List[str]],
enable_key: Optional[Union[str, int]] = None,
enable: bool = True,
kpt_thr: float = 0.5):
super().__init__(
name=name,
input_buffer=input_buffer,
output_buffer=output_buffer,
enable_key=enable_key,
enable=enable)
self.kpt_thr = kpt_thr
def draw(self, input_msg):
canvas = input_msg.get_image()
objects = input_msg.get_objects(lambda x:
('keypoints' in x and 'bbox' in x))
for dataset_meta, group in groupby(objects,
lambda x: x['dataset_meta']):
left_eye_index, right_eye_index = get_eye_keypoint_ids(
dataset_meta)
canvas = self.apply_bigeye_effect(canvas, group, left_eye_index,
right_eye_index)
return canvas
def apply_bigeye_effect(self, canvas: np.ndarray, objects: List[Dict],
left_eye_index: int,
right_eye_index: int) -> np.ndarray:
"""Apply big-eye effect.
Args:
canvas (np.ndarray): The image to apply the effect
objects (list[dict]): The object list with bbox and keypoints
- "bbox" ([K, 4(or 5)]): bbox in [x1, y1, x2, y2, (score)]
- "keypoints" ([K,3]): keypoints in [x, y, score]
left_eye_index (int): Keypoint index of left eye
right_eye_index (int): Keypoint index of right eye
Returns:
np.ndarray: Processed image.
"""
xx, yy = np.meshgrid(
np.arange(canvas.shape[1]), np.arange(canvas.shape[0]))
xx = xx.astype(np.float32)
yy = yy.astype(np.float32)
for obj in objects:
bbox = obj['bbox']
kpts = obj['keypoints']
kpt_scores = obj['keypoint_scores']
if kpt_scores[left_eye_index] < self.kpt_thr or kpt_scores[
right_eye_index] < self.kpt_thr:
continue
kpt_leye = kpts[left_eye_index, :2]
kpt_reye = kpts[right_eye_index, :2]
for xc, yc in [kpt_leye, kpt_reye]:
# distortion parameters
k1 = 0.001
epe = 1e-5
scale = (bbox[2] - bbox[0])**2 + (bbox[3] - bbox[1])**2
r2 = ((xx - xc)**2 + (yy - yc)**2)
r2 = (r2 + epe) / scale # normalized by bbox scale
xx = (xx - xc) / (1 + k1 / r2) + xc
yy = (yy - yc) / (1 + k1 / r2) + yc
canvas = cv2.remap(
canvas,
xx,
yy,
interpolation=cv2.INTER_AREA,
borderMode=cv2.BORDER_REPLICATE)
return canvas
|