File size: 5,755 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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
# 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, load_image_from_disk_or_url
from ..base_visualizer_node import BaseVisualizerNode
from ..registry import NODES


@NODES.register_module()
class SunglassesEffectNode(BaseVisualizerNode):
    """Apply sunglasses effect (draw sunglasses at the facial area)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-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 score threshold of valid keypoints. Default: 0.5
        resource_img_path (str, optional): The resource image path or url.
            The image should be a pair of sunglasses with white background.
            If not specified, the url of a default image will be used. See
            ``SunglassesNode.default_resource_img_path``. Default: ``None``

    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)
    """

    # The image attributes to:
    # "https://www.vecteezy.com/vector-art/1932353-summer-sunglasses-
    # accessory-isolated-icon" by Vecteezy
    default_resource_img_path = (
        'https://user-images.githubusercontent.com/15977946/'
        '170850839-acc59e26-c6b3-48c9-a9ec-87556edb99ed.jpg')

    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,
                 resource_img_path: Optional[str] = None):

        super().__init__(
            name=name,
            input_buffer=input_buffer,
            output_buffer=output_buffer,
            enable_key=enable_key,
            enable=enable)

        if resource_img_path is None:
            resource_img_path = self.default_resource_img_path

        self.resource_img = load_image_from_disk_or_url(resource_img_path)
        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)

        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_sunglasses_effect(canvas, group,
                                                  left_eye_index,
                                                  right_eye_index)
        return canvas

    def apply_sunglasses_effect(self, canvas: np.ndarray, objects: List[Dict],
                                left_eye_index: int,
                                right_eye_index: int) -> np.ndarray:
        """Apply sunglasses effect.

        Args:
            canvas (np.ndarray): The image to apply the effect
            objects (list[dict]): The object list with keypoints
                - "keypoints" ([K,3]): keypoints in [x, y, score]
            left_eye_index (int): Keypoint index of the left eye
            right_eye_index (int): Keypoint index of the right eye

        Returns:
            np.ndarray: Processed image
        """

        hm, wm = self.resource_img.shape[:2]
        # anchor points in the sunglasses image
        pts_src = np.array([[0.3 * wm, 0.3 * hm], [0.3 * wm, 0.7 * hm],
                            [0.7 * wm, 0.3 * hm], [0.7 * wm, 0.7 * hm]],
                           dtype=np.float32)

        for obj in objects:
            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]
            # orthogonal vector to the left-to-right eyes
            vo = 0.5 * (kpt_reye - kpt_leye)[::-1] * [-1, 1]

            # anchor points in the image by eye positions
            pts_tar = np.vstack(
                [kpt_reye + vo, kpt_reye - vo, kpt_leye + vo, kpt_leye - vo])

            h_mat, _ = cv2.findHomography(pts_src, pts_tar)
            patch = cv2.warpPerspective(
                self.resource_img,
                h_mat,
                dsize=(canvas.shape[1], canvas.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)
            canvas = cv2.copyTo(patch, mask, canvas)

        return canvas