diff --git a/common/polystar/common/view/blend_object_on_image.py b/common/polystar/common/view/blend_object_on_image.py
deleted file mode 100644
index 7d8aad3456cf3131069e6942b91188e9d7440e45..0000000000000000000000000000000000000000
--- a/common/polystar/common/view/blend_object_on_image.py
+++ /dev/null
@@ -1,26 +0,0 @@
-import cv2
-
-from polystar.common.models.image import Image
-from polystar.common.models.object import Object
-from polystar.common.view.blend_text_on_image import blend_boxed_text_on_image
-
-_COLORS = [
- [31, 119, 180],
- [255, 127, 14],
- [44, 160, 44],
- [214, 39, 40],
- [148, 103, 189],
- [140, 86, 75],
- [227, 119, 194],
- [127, 127, 127],
- [188, 189, 34],
- [23, 190, 207],
-] # seaborn.color_palette() * 255
-
-
-def blend_object_on_image(image: Image, obj: Object):
- color = _COLORS[obj.type.value]
- cv2.rectangle(image, (obj.x, obj.y), (obj.x + obj.w, obj.y + obj.h), color, 2)
-
- blend_boxed_text_on_image(image, f"{obj.type.name} ({obj.confidence:.1%})", (obj.x, obj.y), _COLORS[obj.type.value])
- return image
diff --git a/common/polystar/common/view/blend_text_on_image.py b/common/polystar/common/view/blend_text_on_image.py
deleted file mode 100644
index 8362a5e4d82df0b05c7bff7b0322682eafff5d7b..0000000000000000000000000000000000000000
--- a/common/polystar/common/view/blend_text_on_image.py
+++ /dev/null
@@ -1,43 +0,0 @@
-from typing import Tuple
-
-import cv2
-import numpy as np
-
-from polystar.common.models.image import Image
-
-ALPHA = 0.5
-FONT = cv2.FONT_HERSHEY_PLAIN
-TEXT_SCALE = 1.0
-TEXT_THICKNESS = 1
-BLACK = (0, 0, 0)
-WHITE = (255, 255, 255)
-
-
-def blend_boxed_text_on_image(img: Image, text: str, topleft: Tuple[int, int], color: Tuple[int, int, int]):
- img_h, img_w, _ = img.shape
- if topleft[0] >= img_w or topleft[1] >= img_h:
- return img
- margin = 3
- size = cv2.getTextSize(text, FONT, TEXT_SCALE, TEXT_THICKNESS)
- w = size[0][0] + margin * 2
- h = size[0][1] + margin * 2
- # the patch is used to draw boxed text
- patch = np.zeros((h, w, 3), dtype=np.uint8)
- patch[...] = color
- cv2.putText(
- patch,
- text,
- (margin + 1, h - margin - 2),
- FONT,
- TEXT_SCALE,
- WHITE,
- thickness=TEXT_THICKNESS,
- lineType=cv2.LINE_8,
- )
- cv2.rectangle(patch, (0, 0), (w - 1, h - 1), BLACK, thickness=1)
- w = min(w, img_w - topleft[0]) # clip overlay at image boundary
- h = min(h, img_h - topleft[1])
- # Overlay the boxed text onto region of interest (roi) in img
- roi = img[topleft[1] : topleft[1] + h, topleft[0] : topleft[0] + w, :]
- cv2.addWeighted(patch[0:h, 0:w, :], ALPHA, roi, 1 - ALPHA, 0, roi)
- return img
diff --git a/common/polystar/common/view/cv2_results_viewer.py b/common/polystar/common/view/cv2_results_viewer.py
new file mode 100644
index 0000000000000000000000000000000000000000..3894982d9c1ee40a93830b744fdd76e99826b222
--- /dev/null
+++ b/common/polystar/common/view/cv2_results_viewer.py
@@ -0,0 +1,73 @@
+import cv2
+import numpy as np
+
+from polystar.common.models.image import Image
+from polystar.common.view.results_viewer_abc import ResultViewerABC, ColorView
+
+ALPHA = 0.5
+FONT = cv2.FONT_HERSHEY_PLAIN
+TEXT_SCALE = 1.0
+TEXT_THICKNESS = 1
+BLACK = (0, 0, 0)
+WHITE = (255, 255, 255)
+COLORS = [
+ (31, 119, 180),
+ (255, 127, 14),
+ (44, 160, 44),
+ (214, 39, 40),
+ (148, 103, 189),
+ (140, 86, 75),
+ (227, 119, 194),
+ (127, 127, 127),
+ (188, 189, 34),
+ (23, 190, 207),
+] # seaborn.color_palette() * 255
+
+
+class CV2ResultViewer(ResultViewerABC):
+ def __init__(self, name: str, delay: int = 1, end_keys: str = "q"):
+ self.end_keys = [ord(c) for c in end_keys]
+ self.delay = delay
+ self.name = name
+ self._current_image: Image = None
+ self.finished = False
+ super().__init__(COLORS)
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ cv2.destroyWindow(self.name)
+
+ def new(self, image: Image):
+ self.height, self.width, _ = image.shape
+ self._current_image = image
+
+ def add_text(self, text: str, x: int, y: int, color: ColorView):
+ margin = 3
+ size = cv2.getTextSize(text, FONT, TEXT_SCALE, TEXT_THICKNESS)
+ w = size[0][0] + margin * 2
+ h = size[0][1] + margin * 2
+ # the patch is used to draw boxed text
+ patch = np.zeros((h, w, 3), dtype=np.uint8)
+ patch[...] = color
+ cv2.putText(
+ patch,
+ text,
+ (margin + 1, h - margin - 2),
+ FONT,
+ TEXT_SCALE,
+ WHITE,
+ thickness=TEXT_THICKNESS,
+ lineType=cv2.LINE_8,
+ )
+ cv2.rectangle(patch, (0, 0), (w - 1, h - 1), BLACK, thickness=1)
+ w = min(w, self.width - x) # clip overlay at image boundary
+ h = min(h, self.height - y)
+ # Overlay the boxed text onto region of interest (roi) in img
+ roi = self._current_image[y : y + h, x : x + w, :]
+ cv2.addWeighted(patch[0:h, 0:w, :], ALPHA, roi, 1 - ALPHA, 0, roi)
+
+ def add_rectangle(self, x: int, y: int, w: int, h: int, color: ColorView):
+ cv2.rectangle(self._current_image, (x, y), (x + w, y + h), color, 2)
+
+ def display(self):
+ cv2.imshow(self.name, self._current_image)
+ self.finished = cv2.waitKey(self.delay) & 0xFF in self.end_keys
diff --git a/common/polystar/common/view/plt_display_image_with_annotation.py b/common/polystar/common/view/plt_display_image_with_annotation.py
deleted file mode 100644
index 8c28081a8bf2569ba08c2ca3b999618bfe34285f..0000000000000000000000000000000000000000
--- a/common/polystar/common/view/plt_display_image_with_annotation.py
+++ /dev/null
@@ -1,29 +0,0 @@
-from typing import Iterable
-
-import matplotlib.pyplot as plt
-
-from polystar.common.models.image import Image
-from polystar.common.models.image_annotation import ImageAnnotation
-from polystar.common.models.object import Object
-import seaborn as sns
-
-
-_COLORS = sns.color_palette()
-
-
-def display_image_with_objects(image: Image, objects: Iterable[Object]):
- plt.figure(figsize=(16, 9))
- plt.imshow(image)
- for obj in objects:
- if obj.confidence >= 0.5:
- color = _COLORS[obj.type.value]
- rect = plt.Rectangle((obj.x, obj.y), obj.w, obj.h, linewidth=2, edgecolor=color, fill=False)
- plt.gca().add_patch(rect)
- plt.text(obj.x, obj.y - 2, f"{obj.type.name} ({int(obj.confidence*100)}%)", color=color, weight="bold")
- plt.axis("off")
- plt.tight_layout()
- plt.show()
-
-
-def display_image_annotation(image_annotation: ImageAnnotation):
- display_image_with_objects(image_annotation.image, image_annotation.objects)
diff --git a/common/polystar/common/view/plt_display_image_with_object.py b/common/polystar/common/view/plt_display_image_with_object.py
deleted file mode 100644
index 075bc85d9812b3c00be55e8a660e39869434cad6..0000000000000000000000000000000000000000
--- a/common/polystar/common/view/plt_display_image_with_object.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from polystar.common.models.image import Image
-from polystar.common.models.object import Object
-from polystar.common.view.plt_display_image_with_annotation import display_image_with_objects
-
-
-def display_object(image: Image, obj: Object):
- display_image_with_objects(image, [obj])
diff --git a/common/polystar/common/view/plt_results_viewer.py b/common/polystar/common/view/plt_results_viewer.py
new file mode 100644
index 0000000000000000000000000000000000000000..e46e019edb1eeccf8bc32da9901f84362c5e5aa5
--- /dev/null
+++ b/common/polystar/common/view/plt_results_viewer.py
@@ -0,0 +1,46 @@
+from typing import Tuple
+
+import matplotlib.pyplot as plt
+
+from polystar.common.models.image import Image
+from polystar.common.view.results_viewer_abc import ResultViewerABC, ColorView
+
+
+COLORS = [
+ (0.12156862745098039, 0.4666666666666667, 0.7058823529411765),
+ (1.0, 0.4980392156862745, 0.054901960784313725),
+ (0.17254901960784313, 0.6274509803921569, 0.17254901960784313),
+ (0.8392156862745098, 0.15294117647058825, 0.1568627450980392),
+ (0.5803921568627451, 0.403921568627451, 0.7411764705882353),
+ (0.5490196078431373, 0.33725490196078434, 0.29411764705882354),
+ (0.8901960784313725, 0.4666666666666667, 0.7607843137254902),
+ (0.4980392156862745, 0.4980392156862745, 0.4980392156862745),
+ (0.7372549019607844, 0.7411764705882353, 0.13333333333333333),
+ (0.09019607843137255, 0.7450980392156863, 0.8117647058823529),
+]
+
+
+class PltResultViewer(ResultViewerABC):
+ def __init__(self, name: str, fig_size: Tuple[int, int] = (16, 9)):
+ self.name = name
+ self.fig_size = fig_size
+ self._current_fig = None
+ super().__init__(COLORS)
+
+ def new(self, image: Image):
+ self._current_fig = plt.figure(figsize=self.fig_size)
+ plt.imshow(image)
+ plt.title(self.name)
+ plt.axis("off")
+
+ def add_text(self, text: str, x: int, y: int, color: ColorView):
+ plt.text(x, y - 2, text, color=color, weight="bold")
+
+ def add_rectangle(self, x: int, y: int, w: int, h: int, color: ColorView):
+ rect = plt.Rectangle((x, y), w, h, linewidth=2, edgecolor=color, fill=False)
+ plt.gca().add_patch(rect)
+
+ def display(self):
+ plt.tight_layout()
+ plt.show()
+ plt.close(self._current_fig)
diff --git a/common/polystar/common/view/results_viewer_abc.py b/common/polystar/common/view/results_viewer_abc.py
new file mode 100644
index 0000000000000000000000000000000000000000..486bf0d29bb288453b158acde5ef0d95e084f974
--- /dev/null
+++ b/common/polystar/common/view/results_viewer_abc.py
@@ -0,0 +1,52 @@
+from abc import ABC, abstractmethod
+from typing import Iterable, Tuple, NewType, Sequence
+
+from polystar.common.models.image import Image
+from polystar.common.models.image_annotation import ImageAnnotation
+from polystar.common.models.object import Object
+
+ColorView = NewType("ColorView", Tuple[float, float, float])
+
+
+class ResultViewerABC(ABC):
+ def __init__(self, colors: Sequence[ColorView]):
+ self.colors = colors
+
+ def __enter__(self):
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ pass
+
+ @abstractmethod
+ def new(self, image: Image):
+ pass
+
+ @abstractmethod
+ def add_text(self, text: str, x: int, y: int, color: ColorView):
+ pass
+
+ @abstractmethod
+ def add_rectangle(self, x: int, y: int, w: int, h: int, color: ColorView):
+ pass
+
+ @abstractmethod
+ def display(self):
+ pass
+
+ def add_object(self, obj: Object):
+ color = self.colors[obj.type.value]
+ self.add_rectangle(obj.x, obj.y, obj.w, obj.h, color)
+ self.add_text(f"{obj.type.name} ({obj.confidence:.1%})", obj.x, obj.y, color)
+
+ def add_objects(self, objects: Iterable[Object]):
+ for obj in objects:
+ self.add_object(obj)
+
+ def display_image_with_objects(self, image: Image, objects: Iterable[Object]):
+ self.new(image)
+ self.add_objects(objects)
+ self.display()
+
+ def display_image_annotation(self, annotation: ImageAnnotation):
+ self.display_image_with_objects(annotation.image, annotation.objects)
diff --git a/robots-at-robots/research/demos/demo_infer.py b/robots-at-robots/research/demos/demo_infer.py
index 4e0a7f237d67847d916a9318e520897fdea3fe2c..9769460c3cf395f0d1729138d17683c3695bbc46 100644
--- a/robots-at-robots/research/demos/demo_infer.py
+++ b/robots-at-robots/research/demos/demo_infer.py
@@ -2,7 +2,7 @@ from polystar.common.models.label_map import LabelMap
from polystar.common.pipeline.objects_detectors.tf_model_objects_detector import TFModelObjectsDetector
from polystar.common.pipeline.objects_validators.confidence_object_validator import ConfidenceObjectValidator
from polystar.common.utils.tensorflow import patch_tf_v2
-from polystar.common.view.plt_display_image_with_annotation import display_image_with_objects
+from polystar.common.view.plt_results_viewer import PltResultViewer
from polystar.robots_at_robots.dependency_injection import make_injector
from research.demos.utils import load_tf_model
from research_common.dataset.dji.dji_roco_datasets import DJIROCODataset
@@ -16,11 +16,13 @@ if __name__ == "__main__":
objects_detector = TFModelObjectsDetector(load_tf_model(), injector.get(LabelMap))
filters = [ConfidenceObjectValidator(confidence_threshold=0.5)]
- for i, image in enumerate(SplitDataset(DJIROCODataset.CentralChina, Split.Test).images):
- objects = objects_detector.detect(image)
- for f in filters:
- objects = f.filter(objects, image)
+ with PltResultViewer("Demo of tf model") as viewer:
+ for i, image in enumerate(SplitDataset(DJIROCODataset.CentralChina, Split.Test).images):
+ objects = objects_detector.detect(image)
+ for f in filters:
+ objects = f.filter(objects, image)
- display_image_with_objects(image, objects)
- if i == 0:
- break
+ viewer.display_image_with_objects(image, objects)
+
+ if i == 5:
+ break
diff --git a/robots-at-robots/research/demos/demo_pipeline.py b/robots-at-robots/research/demos/demo_pipeline.py
index 047aaf2a7d0c120fb5f6dd9e81ffb3636e43a918..ca58ed5fbc583e81a082ff747036caa7ea3a4a00 100644
--- a/robots-at-robots/research/demos/demo_pipeline.py
+++ b/robots-at-robots/research/demos/demo_pipeline.py
@@ -10,7 +10,7 @@ from polystar.common.pipeline.objects_validators.type_object_validator import Ty
from polystar.common.pipeline.pipeline import Pipeline
from polystar.common.pipeline.target_factories.ratio_simple_target_factory import RatioSimpleTargetFactory
from polystar.common.utils.tensorflow import patch_tf_v2
-from polystar.common.view.plt_display_image_with_object import display_object
+from polystar.common.view.plt_results_viewer import PltResultViewer
from polystar.robots_at_robots.dependency_injection import make_injector
from research.demos.utils import load_tf_model
from research_common.dataset.dji.dji_roco_datasets import DJIROCODataset
@@ -28,11 +28,12 @@ if __name__ == "__main__":
target_factory=RatioSimpleTargetFactory(injector.get(Camera), 300, 100),
)
- for i, image_path in enumerate(SplitDataset(DJIROCODataset.CentralChina, Split.Test).image_paths):
- image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
- obj = pipeline.predict_best_object(image)
+ with PltResultViewer("Demo of tf model") as viewer:
+ for i, image_path in enumerate(SplitDataset(DJIROCODataset.CentralChina, Split.Test).image_paths):
+ image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
+ obj = pipeline.predict_best_object(image)
- display_object(image, obj)
+ viewer.display_image_with_objects(image, [obj])
- if i == 0:
- break
+ if i == 5:
+ break
diff --git a/robots-at-robots/research/demos/demo_pipeline_camera.py b/robots-at-robots/research/demos/demo_pipeline_camera.py
index ffe25067e4d163eb3c2bdc33075b8efd2a92afe3..33f15701452d229d2e76b445bf0ed79d507964db 100644
--- a/robots-at-robots/research/demos/demo_pipeline_camera.py
+++ b/robots-at-robots/research/demos/demo_pipeline_camera.py
@@ -1,6 +1,5 @@
from time import time
-import cv2
import pycuda.autoinit # This is needed for initializing CUDA driver
from polystar.common.constants import MODELS_DIR
@@ -10,14 +9,10 @@ from polystar.common.models.trt_model import TRTModel
from polystar.common.pipeline.objects_detectors.trt_model_object_detector import TRTModelObjectsDetector
from polystar.common.pipeline.objects_validators.confidence_object_validator import ConfidenceObjectValidator
from polystar.common.utils.tensorflow import patch_tf_v2
-from polystar.common.view.blend_object_on_image import blend_object_on_image
-from polystar.common.view.blend_text_on_image import blend_boxed_text_on_image
+from polystar.common.view.cv2_results_viewer import CV2ResultViewer
from polystar.robots_at_robots.dependency_injection import make_injector
from polystar.robots_at_robots.globals import settings
-WINDOWS_NAME = "TensorRT demo"
-
-
[pycuda.autoinit] # So pycharm won't remove the import
@@ -31,23 +26,21 @@ if __name__ == "__main__":
filters = [ConfidenceObjectValidator(confidence_threshold=0.5)]
fps = 0
- try:
+ with CV2ResultViewer("TensorRT demo") as viewer:
for image in CameraFrameGenerator(1_280, 720).generate():
+
previous_time = time()
+
+ # inference
objects = objects_detector.detect(image)
for f in filters:
objects = f.filter(objects, image)
+ # display
fps = 0.9 * fps + 0.1 / (time() - previous_time)
- blend_boxed_text_on_image(image, f"FPS: {fps:.1f}", (10, 10), (0, 0, 0))
-
- for obj in objects:
- blend_object_on_image(image, obj)
-
- # Display the resulting frame
- cv2.imshow("frame", image)
- if cv2.waitKey(1) & 0xFF == ord("q"):
+ viewer.new(image)
+ viewer.add_objects(objects)
+ viewer.add_text(f"FPS: {fps:.1f}", 10, 10, (0, 0, 0))
+ viewer.display()
+ if viewer.finished:
break
- finally:
- # When everything done, release the capture
- cv2.destroyAllWindows()