def __init__(self,
name: str,
enable_key: Optional[Union[str, int]] = None,
max_fps: int = 30,
input_check_interval: float = 0.01,
enable: bool = True,
daemon: bool = False,
multi_input: bool = False):
super().__init__(name=name, daemon=daemon)
self._executor = None
self._enabled = enable
self.enable_key = enable_key
self.max_fps = max_fps
self.input_check_interval = input_check_interval
self.multi_input = multi_input
# A partitioned buffer manager the executor'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 executor'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_executor() 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,
)
# Logger
self.logger = logging.getLogger(f'Node "{self.name}"')
def inference_pose():
print('Thread "pose" started')
stop_watch = StopWatch(window=10)
while True:
while len(det_result_queue) < 1:
time.sleep(0.001)
with det_result_queue_mutex:
ts_input, frame, t_info, mmdet_results = det_result_queue.popleft()
pose_results_list = []
for model_info, pose_history in zip(pose_model_list,
pose_history_list):
model_name = model_info['name']
pose_model = model_info['model']
cat_ids = model_info['cat_ids']
pose_results_last = pose_history['pose_results_last']
next_id = pose_history['next_id']
with stop_watch.timeit(model_name):
# process mmdet results
det_results = process_mmdet_results(
mmdet_results,
class_names=det_model.CLASSES,
cat_ids=cat_ids)
# inference pose model
dataset_name = pose_model.cfg.data['test']['type']
pose_results, _ = inference_top_down_pose_model(
pose_model,
frame,
det_results,
bbox_thr=args.det_score_thr,
format='xyxy',
dataset=dataset_name)
pose_results, next_id = get_track_id(pose_results,
pose_results_last,
next_id,
use_oks=False,
tracking_thr=0.3,
use_one_euro=True,
fps=None)
pose_results_list.append(pose_results)
# update pose history
pose_history['pose_results_last'] = pose_results
pose_history['next_id'] = next_id
t_info += stop_watch.report_strings()
with pose_result_queue_mutex:
pose_result_queue.append((ts_input, t_info, pose_results_list))
event_inference_done.set()
def inference_detection():
print('Thread "det" started')
stop_watch = StopWatch(window=10)
min_interval = 1.0 / args.inference_fps
_ts_last = None # timestamp when last inference was done
while True:
while len(input_queue) < 1:
time.sleep(0.001)
with input_queue_mutex:
ts_input, frame = input_queue.popleft()
# inference detection
with stop_watch.timeit('Det'):
mmdet_results = inference_detector(det_model, frame)
t_info = stop_watch.report_strings()
with det_result_queue_mutex:
det_result_queue.append((ts_input, frame, t_info, mmdet_results))
# limit the inference FPS
_ts = time.time()
if _ts_last is not None and _ts - _ts_last < min_interval:
time.sleep(min_interval - _ts + _ts_last)
_ts_last = time.time()
def test_stopwatch():
window_size = 5
test_loop = 10
outer_time = 100
inner_time = 100
stop_watch = StopWatch(window=window_size)
for _ in range(test_loop):
with stop_watch.timeit():
time.sleep(outer_time / 1000.)
with stop_watch.timeit('inner'):
time.sleep(inner_time / 1000.)
_ = stop_watch.report()
_ = stop_watch.report_strings()
def test_stopwatch():
window_size = 5
test_loop = 10
outer_time = 100
inner_time = 100
stop_watch = StopWatch(window=window_size)
for _ in range(test_loop):
with stop_watch.timeit():
time.sleep(outer_time / 1000.)
with stop_watch.timeit('inner'):
time.sleep(inner_time / 1000.)
report = stop_watch.report()
_ = stop_watch.report_strings()
np.testing.assert_allclose(report['_FPS_'],
outer_time + inner_time,
rtol=0.01)
np.testing.assert_allclose(report['inner'], inner_time, rtol=0.01)
def display():
print('Thread "display" started')
stop_watch = StopWatch(window=10)
# initialize result status
ts_inference = None # timestamp of the latest inference result
fps_inference = 0. # infenrece FPS
t_delay_inference = 0. # inference result time delay
pose_results_list = None # latest inference result
t_info = [] # upstream time information (list[str])
# initialize visualization and output
sunglasses_img = None # resource image for sunglasses effect
text_color = (228, 183, 61) # text color to show time/system information
vid_out = None # video writer
# show instructions
print('Keyboard shortcuts: ')
print('"v": Toggle the visualization of bounding boxes and poses.')
print('"s": Toggle the sunglasses effect.')
print('"b": Toggle the bug-eye effect.')
print('"Q", "q" or Esc: Exit.')
while True:
with stop_watch.timeit('_FPS_'):
# acquire a frame from buffer
ts_input, frame = frame_buffer.get()
# input ending signal
if ts_input is None:
break
img = frame
# get pose estimation results
if len(pose_result_queue) > 0:
with pose_result_queue_mutex:
_result = pose_result_queue.popleft()
_ts_input, t_info, pose_results_list = _result
_ts = time.time()
if ts_inference is not None:
fps_inference = 1.0 / (_ts - ts_inference)
ts_inference = _ts
t_delay_inference = (_ts - _ts_input) * 1000
# visualize detection and pose results
if pose_results_list is not None:
for model_info, pose_results in zip(pose_model_list,
pose_results_list):
pose_model = model_info['model']
bbox_color = model_info['bbox_color']
dataset_name = pose_model.cfg.data['test']['type']
# show pose results
if args.show_pose:
img = vis_pose_result(pose_model,
img,
pose_results,
radius=4,
thickness=2,
dataset=dataset_name,
kpt_score_thr=args.kpt_thr,
bbox_color=bbox_color)
# sunglasses effect
if args.sunglasses:
if dataset_name == 'TopDownCocoDataset':
left_eye_idx = 1
right_eye_idx = 2
elif dataset_name == 'AnimalPoseDataset':
left_eye_idx = 0
right_eye_idx = 1
else:
raise ValueError(
'Sunglasses effect does not support'
f'{dataset_name}')
if sunglasses_img is None:
# The image attributes to:
# https://www.vecteezy.com/free-vector/glass
# Glass Vectors by Vecteezy
sunglasses_img = cv2.imread(
'demo/resources/sunglasses.jpg')
img = apply_sunglasses_effect(img, pose_results,
sunglasses_img,
left_eye_idx,
right_eye_idx)
# bug-eye effect
if args.bugeye:
if dataset_name == 'TopDownCocoDataset':
left_eye_idx = 1
right_eye_idx = 2
elif dataset_name == 'AnimalPoseDataset':
left_eye_idx = 0
right_eye_idx = 1
else:
raise ValueError('Bug-eye effect does not support'
f'{dataset_name}')
img = apply_bugeye_effect(img, pose_results,
left_eye_idx, right_eye_idx)
# delay control
if args.display_delay > 0:
t_sleep = args.display_delay * 0.001 - (time.time() - ts_input)
if t_sleep > 0:
time.sleep(t_sleep)
t_delay = (time.time() - ts_input) * 1000
# show time information
t_info_display = stop_watch.report_strings() # display fps
t_info_display.append(f'Inference FPS: {fps_inference:>5.1f}')
t_info_display.append(f'Delay: {t_delay:>3.0f}')
t_info_display.append(
f'Inference Delay: {t_delay_inference:>3.0f}')
t_info_str = ' | '.join(t_info_display + t_info)
cv2.putText(img, t_info_str, (20, 20), cv2.FONT_HERSHEY_DUPLEX,
0.3, text_color, 1)
# collect system information
sys_info = [
f'RES: {img.shape[1]}x{img.shape[0]}',
f'Buffer: {frame_buffer.qsize()}/{frame_buffer.maxsize}'
]
if psutil_proc is not None:
sys_info += [
f'CPU: {psutil_proc.cpu_percent():.1f}%',
f'MEM: {psutil_proc.memory_percent():.1f}%'
]
sys_info_str = ' | '.join(sys_info)
cv2.putText(img, sys_info_str, (20, 40), cv2.FONT_HERSHEY_DUPLEX,
0.3, text_color, 1)
# save the output video frame
if args.out_video_file is not None:
if vid_out is None:
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
fps = args.out_video_fps
frame_size = (img.shape[1], img.shape[0])
vid_out = cv2.VideoWriter(args.out_video_file, fourcc, fps,
frame_size)
vid_out.write(img)
# display
cv2.imshow('mmpose webcam demo', img)
keyboard_input = cv2.waitKey(1)
if keyboard_input in (27, ord('q'), ord('Q')):
break
elif keyboard_input == ord('s'):
args.sunglasses = not args.sunglasses
elif keyboard_input == ord('b'):
args.bugeye = not args.bugeye
elif keyboard_input == ord('v'):
args.show_pose = not args.show_pose
cv2.destroyAllWindows()
if vid_out is not None:
vid_out.release()
event_exit.set()
class Node(Thread, metaclass=ABCMeta):
"""Base class for node, which is the interface of basic function module.
:class:`Node` inherits :class:`threading.Thread`. All subclasses should
override following methods:
- ``process()``
- ``bypass()`` (optional)
Parameters:
name (str): 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-keys are reserved for
particular use. For example: 'q', 'Q' and 27 are used for exiting.
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``
multi_input (bool): Whether load all messages in buffer. If False,
only one message will be loaded each time. Default: ``False``
"""
def __init__(self,
name: str,
enable_key: Optional[Union[str, int]] = None,
max_fps: int = 30,
input_check_interval: float = 0.01,
enable: bool = True,
daemon: bool = False,
multi_input: bool = False):
super().__init__(name=name, daemon=daemon)
self._executor = None
self._enabled = enable
self.enable_key = enable_key
self.max_fps = max_fps
self.input_check_interval = input_check_interval
self.multi_input = multi_input
# A partitioned buffer manager the executor'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 executor'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_executor() 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,
)
# Logger
self.logger = logging.getLogger(f'Node "{self.name}"')
@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,
trigger: 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.
The subclass of Node should invoke `register_input_buffer` in its
`__init__` method. This method can be invoked multiple times to
register multiple input buffers.
Args:
buffer_name (str): The name of the buffer
input_name (str): The name of the fetched message from the
corresponding buffer
trigger (bool): An trigger 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, trigger)
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_executor(self, executor):
"""Assign the node to an executor so the node can access the buffers
and event manager of the executor.
This method should be invoked by the executor instance.
Args:
executor (:obj:`WebcamExecutor`): The executor to hold the node
"""
# Get partitioned buffer manager
buffer_names = [
buffer.buffer_name
for buffer in self._input_buffers + self._output_buffers
]
self._buffer_manager = executor.buffer_manager.get_sub_manager(
buffer_names)
# Get event manager
self._event_manager = executor.event_manager
def _get_input_from_buffer(self) -> Tuple[bool, Optional[Dict]]:
"""Get and pack input data.
The function returns a tuple (status, data). If the trigger buffers
are ready, the status flag will be True, and the packed data is a dict
whose items are buffer names and corresponding messages (unready
non-trigger buffers will give a `None`). Otherwise, the status flag is
False and the packed data is None.
Returns:
tuple[bool, dict]: The first item is a bool value indicating
whether input is ready (i.e., all tirgger buffers are ready). The
second value is a dict of buffer names and messages.
"""
buffer_manager = self._buffer_manager
if buffer_manager is None:
raise ValueError(f'Node "{self.name}": not set to an executor.')
# Check that trigger buffers are ready
for buffer_info in self._input_buffers:
if buffer_info.trigger 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:
while not buffer_manager.is_empty(buffer_info.buffer_name):
msg = buffer_manager.get(buffer_info.buffer_name, block=False)
if self.multi_input:
if result[buffer_info.input_name] is None:
result[buffer_info.input_name] = []
result[buffer_info.input_name].append(msg)
else:
result[buffer_info.input_name] = msg
break
# Return unsuccessful flag if any trigger input is unready
if buffer_info.trigger and result[buffer_info.input_name] is None:
return False, None
return True, result
def _send_output_to_buffers(self, output_msg):
"""Send output of ``process()`` to the registered output buffers.
Args:
output_msg (Message): output message
"""
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 method that implements 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[str, :obj:`Message`]): The input data collected
from the buffers. For each item, the key is the `input_name`
of the registered input buffer, and the value is a Message
instance fetched from the buffer (or None if the buffer is
non-trigger and not ready).
Returns:
Message: The output message of the node which 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 a node must override this method if it has `enable_key`.
This method has the same signature as ``process()``.
Args:
input_msgs (dict[str, :obj:`Message`]): The input data collected
from the buffers. For each item, the key is the `input_name`
of the registered input buffer, and the value is a Message
instance fetched from the buffer (or None if the buffer is
non-trigger and not ready).
Returns:
Message: The output message of the node which will be send to all
registered output buffers.
"""
raise NotImplementedError
def _get_node_info(self) -> Dict:
"""Get route information of the node.
Default information includes:
- ``'fps'``: The processing speed of the node
- ``'timestamp'``: The time that this method is invoked
Subclasses can override this method to customize the node information.
Returns:
dict: The items of node information
"""
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.
Subclasses of :class:`Node` should not override this method in
subclasses.
"""
self.logger.info('Process 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)
self.logger.info('Process ends.')