Skip to content
Snippets Groups Projects
Commit 92bda7e7 authored by Mathieu Beligon's avatar Mathieu Beligon
Browse files

[CV2ResultsViewer] Add callback support

parent 48bbd821
No related branches found
No related tags found
No related merge requests found
from contextlib import suppress
from typing import Any, Callable, Dict
import cv2 import cv2
import numpy as np import numpy as np
...@@ -23,10 +26,12 @@ COLORS = [ ...@@ -23,10 +26,12 @@ COLORS = [
(23, 190, 207), (23, 190, 207),
] # seaborn.color_palette() * 255 ] # seaborn.color_palette() * 255
Callback = Callable[[], Any]
class CV2ResultViewer(ResultViewerABC): class CV2ResultViewer(ResultViewerABC):
def __init__(self, name: str, delay: int = 1, end_keys: str = "q"): def __init__(self, name: str, delay: int = 1, end_key: str = "q", key_callbacks: Dict[str, Callback] = None):
self.end_keys = [ord(c) for c in end_keys] self.keycode_callbacks = self._make_keycode_callbacks(end_key, key_callbacks or {})
self.delay = delay self.delay = delay
self.name = name self.name = name
self._current_image: Image = None self._current_image: Image = None
...@@ -70,4 +75,13 @@ class CV2ResultViewer(ResultViewerABC): ...@@ -70,4 +75,13 @@ class CV2ResultViewer(ResultViewerABC):
def display(self): def display(self):
cv2.imshow(self.name, self._current_image) cv2.imshow(self.name, self._current_image)
self.finished = cv2.waitKey(self.delay) & 0xFF in self.end_keys keycode = cv2.waitKey(self.delay) & 0xFF
with suppress(KeyError):
self.keycode_callbacks[keycode]()
def stop(self):
self.finished = True
def _make_keycode_callbacks(self, end_key: str, key_callbacks: Dict[str, Callback]) -> Dict[int, Callback]:
key_callbacks[end_key] = self.stop
return {ord(k): f for k, f in key_callbacks.items()}
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment