File size: 5,401 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
# Copyright (c) OpenMMLab. All rights reserved.
import time
from typing import List, Union

from mmpose.utils.timer import RunningAverage
from ..node import Node
from ..registry import NODES


@NODES.register_module()
class ObjectAssignerNode(Node):
    """Assign the object information to the frame message.

    :class:`ObjectAssignerNode` enables asynchronous processing of model
    inference and video I/O, so the video will be captured and displayed
    smoothly regardless of the model inference speed. Specifically,
    :class:`ObjectAssignerNode` takes messages from both model branch and
    video I/O branch as its input, indicated as "object message" and "frame
    message" respectively. When an object message arrives it will update the
    latest object information; and when a frame message arrives, it will be
    assigned with the latest object information and output.

    Specially, if the webcam executor is set to synchrounous mode, the
    behavior of :class:`ObjectAssignerNode` will be different: When an object
    message arrives, it will trigger an output of itself; and the frame
    messages will be ignored.

    Args:
        name (str): The node name (also thread name)
        frame_buffer (str): Buffer name for frame messages
        object_buffer (str): Buffer name for object messages
        output_buffer (str): The name(s) of the output buffer(s)

    Example::
        >>> cfg =dict(
        ...     type='ObjectAssignerNode',
        ...     name='object assigner',
        ...     frame_buffer='_frame_',
        ...     # `_frame_` is an executor-reserved buffer
        ...     object_buffer='animal_pose',
        ...     output_buffer='frame')

        >>> from mmpose.apis.webcam.nodes import NODES
        >>> node = NODES.build(cfg)
    """

    def __init__(self, name: str, frame_buffer: str, object_buffer: str,
                 output_buffer: Union[str, List[str]]):
        super().__init__(name=name, enable=True)
        self.synchronous = None

        # Cache the latest model result
        self.last_object_msg = None
        self.last_output_msg = None

        # Inference speed analysis
        self.frame_fps = RunningAverage(window=10)
        self.frame_lag = RunningAverage(window=10)
        self.object_fps = RunningAverage(window=10)
        self.object_lag = RunningAverage(window=10)

        # Register buffers
        # The trigger buffer depends on the executor.synchronous attribute,
        # so it will be set later after the executor is assigned in
        # ``set_executor``.
        self.register_input_buffer(object_buffer, 'object', trigger=False)
        self.register_input_buffer(frame_buffer, 'frame', trigger=False)
        self.register_output_buffer(output_buffer)

    def set_executor(self, executor):
        super().set_executor(executor)
        # Set synchronous according to the executor
        if executor.synchronous:
            self.synchronous = True
            trigger = 'object'
        else:
            self.synchronous = False
            trigger = 'frame'

        # Set trigger input buffer according to the synchronous setting
        for buffer_info in self._input_buffers:
            if buffer_info.input_name == trigger:
                buffer_info.trigger = True

    def process(self, input_msgs):
        object_msg = input_msgs['object']

        # Update last result
        if object_msg is not None:
            # Update result FPS
            if self.last_object_msg is not None:
                self.object_fps.update(
                    1.0 /
                    (object_msg.timestamp - self.last_object_msg.timestamp))
            # Update inference latency
            self.object_lag.update(time.time() - object_msg.timestamp)
            # Update last inference result
            self.last_object_msg = object_msg

        if not self.synchronous:
            # Asynchronous mode:
            # Assign the latest object information to the
            # current frame.
            frame_msg = input_msgs['frame']

            self.frame_lag.update(time.time() - frame_msg.timestamp)

            # Assign objects to frame
            if self.last_object_msg is not None:
                frame_msg.update_objects(self.last_object_msg.get_objects())
                frame_msg.merge_route_info(
                    self.last_object_msg.get_route_info())

            output_msg = frame_msg

        else:
            # Synchronous mode:
            # The current frame will be ignored. Instead,
            # the frame from which the latest object information is obtained
            # will be used.
            self.frame_lag.update(time.time() - object_msg.timestamp)
            output_msg = object_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['object_fps'] = self.object_fps.average()
        info['object_lag (ms)'] = self.object_lag.average() * 1000
        info['frame_fps'] = self.frame_fps.average()
        info['frame_lag (ms)'] = self.frame_lag.average() * 1000
        return info