From acc57a9dc778329a67b81a603a60e97efc0ae210 Mon Sep 17 00:00:00 2001
From: Mathieu Beligon <mathieu.beligon@gmail.com>
Date: Fri, 20 Mar 2020 18:35:40 -0400
Subject: [PATCH] [robots] (demo trt) Add target sender to pipeline
---
.../research/demos/demo_pipeline_camera.py | 31 +++++++++++--------
1 file changed, 18 insertions(+), 13 deletions(-)
diff --git a/robots-at-robots/research/demos/demo_pipeline_camera.py b/robots-at-robots/research/demos/demo_pipeline_camera.py
index 8088138..b340ec2 100644
--- a/robots-at-robots/research/demos/demo_pipeline_camera.py
+++ b/robots-at-robots/research/demos/demo_pipeline_camera.py
@@ -4,6 +4,7 @@ from time import time
import pycuda.autoinit # This is needed for initializing CUDA driver
from polystar.common.communication.file_descriptor_target_sender import FileDescriptorTargetSender
+from polystar.common.communication.print_target_sender import PrintTargetSender
from polystar.common.constants import MODELS_DIR
from polystar.common.frame_generators.camera_frame_generator import CameraFrameGenerator
from polystar.common.models.camera import Camera
@@ -15,6 +16,7 @@ from polystar.common.pipeline.object_selectors.closest_object_selector import Cl
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.pipeline.objects_validators.type_object_validator import TypeObjectValidator
+from polystar.common.pipeline.pipeline import NoTargetFound
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.cv2_results_viewer import CV2ResultViewer
@@ -41,20 +43,23 @@ if __name__ == "__main__":
fps = 0
with CV2ResultViewer("TensorRT demo") as viewer:
- for image in CameraFrameGenerator(1_280, 720).generate():
+ for image in CameraFrameGenerator(1_280, 720).generate():
+ viewer.new(image)
previous_time = time()
+ try:
+ # inference
+ pipeline.predict_target(image)
- # inference
- pipeline.predict_target(image)
-
- # display
- fps = 0.9 * fps + 0.1 / (time() - previous_time)
- viewer.new(image)
- viewer.add_objects(pipeline.debug_info_.validated_objects, forced_color=(0.6, 0.6, 0.6))
- viewer.add_object(pipeline.debug_info_.selected_object)
- viewer.add_text(f"FPS: {fps:.1f}", 10, 10, (0, 0, 0))
- viewer.display()
+ # display
+ viewer.add_objects(pipeline.debug_info_.validated_objects, forced_color=(0.6, 0.6, 0.6))
+ viewer.add_object(pipeline.debug_info_.selected_object)
- if viewer.finished:
- break
+ if viewer.finished:
+ break
+ except NoTargetFound:
+ pass
+ finally:
+ fps = 0.9 * fps + 0.1 / (time() - previous_time)
+ viewer.add_text(f"FPS: {fps:.1f}", 10, 10, (0, 0, 0))
+ viewer.display()
--
GitLab