From 951900b0466804ba309214ae4d992e3a8c49aa1d Mon Sep 17 00:00:00 2001
From: Mathieu Beligon <beligonmathieu@gmail.com>
Date: Sun, 7 Mar 2021 16:58:36 -0500
Subject: [PATCH] [TRTObjectDetection] Fix multi threading
---
src/polystar/dependency_injection.py | 11 +++--
.../cv2_frame_generator_abc.py | 3 +-
.../objects_detectors/robots_detector.py | 31 +++++++++-----
.../trt_model_object_detector.py | 41 +++++++++----------
src/research/scripts/demo_pipeline_camera.py | 1 +
5 files changed, 48 insertions(+), 39 deletions(-)
diff --git a/src/polystar/dependency_injection.py b/src/polystar/dependency_injection.py
index 748cb1e..71276be 100644
--- a/src/polystar/dependency_injection.py
+++ b/src/polystar/dependency_injection.py
@@ -1,5 +1,5 @@
from dataclasses import dataclass
-from typing import List
+from typing import List, Type
from injector import Injector, Module, multiprovider, provider, singleton
from numpy.core._multiarray_umath import deg2rad
@@ -25,7 +25,6 @@ from polystar.target_pipeline.objects_linker.objects_linker_abs import ObjectsLi
from polystar.target_pipeline.objects_linker.simple_objects_linker import SimpleObjectsLinker
from polystar.target_pipeline.target_factories.ratio_simple_target_factory import RatioSimpleTargetFactory
from polystar.target_pipeline.target_factories.target_factory_abc import TargetFactoryABC
-from research.common.constants import PIPELINES_DIR
def make_injector() -> Injector:
@@ -57,15 +56,14 @@ class CommonModule(Module):
@provider
@singleton
- def provide_objects_detector(self) -> ObjectsDetectorABC:
- model_dir = PIPELINES_DIR / "roco-detection" / settings.OBJECTS_DETECTION_MODEL
+ def provide_objects_detector_class(self) -> Type[ObjectsDetectorABC]:
if self.settings.is_dev:
from polystar.target_pipeline.objects_detectors.tf_model_objects_detector import TFModelObjectsDetector
- return TFModelObjectsDetector(model_dir)
+ return TFModelObjectsDetector
from polystar.target_pipeline.objects_detectors.trt_model_object_detector import TRTModelObjectsDetector
- return TRTModelObjectsDetector(model_dir)
+ return TRTModelObjectsDetector
@multiprovider
@singleton
@@ -104,6 +102,7 @@ class CommonModule(Module):
return SimpleObjectsLinker(min_percentage_intersection=0.8)
@provider
+ @singleton
def provide_webcam(self) -> FrameGeneratorABC:
if self.settings.is_prod:
return make_csi_camera_frame_generator(1_280, 720)
diff --git a/src/polystar/frame_generators/cv2_frame_generator_abc.py b/src/polystar/frame_generators/cv2_frame_generator_abc.py
index f191ea1..46d7969 100644
--- a/src/polystar/frame_generators/cv2_frame_generator_abc.py
+++ b/src/polystar/frame_generators/cv2_frame_generator_abc.py
@@ -18,7 +18,8 @@ class CV2Capture(Iterator[Image]):
def __init__(self, capture_params: Iterable):
self._cap = VideoCapture(*capture_params)
assert self._cap.isOpened()
- self._cap.set(CAP_PROP_BUFFERSIZE, 0)
+
+ # self._cap.set(CAP_PROP_BUFFERSIZE, 0)
def __next__(self) -> Image:
success, frame = self._cap.read()
diff --git a/src/polystar/target_pipeline/objects_detectors/robots_detector.py b/src/polystar/target_pipeline/objects_detectors/robots_detector.py
index 30e5631..ab69fda 100644
--- a/src/polystar/target_pipeline/objects_detectors/robots_detector.py
+++ b/src/polystar/target_pipeline/objects_detectors/robots_detector.py
@@ -1,6 +1,6 @@
from dataclasses import dataclass
from threading import Condition
-from typing import Iterable, List, Tuple
+from typing import Iterable, List, Tuple, Type
from injector import inject
@@ -8,28 +8,29 @@ from polystar.models.box import Box
from polystar.models.image import Image
from polystar.models.label_map import LabelMap
from polystar.models.roco_object import ObjectType
+from polystar.settings import settings
from polystar.target_pipeline.detected_objects.detected_armor import DetectedArmor
from polystar.target_pipeline.detected_objects.detected_robot import DetectedRobot
from polystar.target_pipeline.detected_objects.objects_params import ObjectParams
from polystar.target_pipeline.objects_detectors.objects_detector_abc import ObjectsDetectorABC
from polystar.target_pipeline.objects_linker.objects_linker_abs import ObjectsLinkerABC
from polystar.utils.thread import MyThread
+from polystar.utils.time import time_it
+from research.common.constants import PIPELINES_DIR
@inject
@dataclass
class RobotsDetector:
label_map: LabelMap
- object_detector: ObjectsDetectorABC
+ object_detector_class: Type[ObjectsDetectorABC]
objects_linker: ObjectsLinkerABC
def flow_robots(self, images: Iterable[Image]) -> Iterable[List[DetectedRobot]]:
- detector_thread = ObjectsDetectorThread(self.object_detector, images)
+ detector_thread = ObjectsDetectorThread(self.object_detector_class, images)
detector_thread.start()
while detector_thread.running:
- with detector_thread.condition:
- detector_thread.condition.wait() # TODO: timeout ?
- image, objects_params = detector_thread.image, detector_thread.objects_params
+ image, objects_params = detector_thread.get_next_detection()
robots, armors = self.make_robots_and_armors(objects_params, image)
yield image, list(self.objects_linker.link_armors_to_robots(robots, armors, image))
@@ -56,16 +57,26 @@ class RobotsDetector:
class ObjectsDetectorThread(MyThread):
- def __init__(self, objects_detector: ObjectsDetectorABC, images: Iterable[Image]):
+ def __init__(self, objects_detector_class: Type[ObjectsDetectorABC], images: Iterable[Image]):
super().__init__()
- self.objects_detector = objects_detector
- self.objects_params: List[ObjectParams] = None
+ self.objects_detector_class = objects_detector_class
self.images = images
self.condition = Condition()
+ self.image = None
+
+ def get_next_detection(self) -> Tuple[Image, List[ObjectParams]]:
+ with self.condition:
+ self.condition.wait()
+ return self.image, self.objects_params
def loop(self):
+ model_dir = PIPELINES_DIR / "roco-detection" / settings.OBJECTS_DETECTION_MODEL
+ objects_detector = self.objects_detector_class(model_dir)
for image in self.images:
- objects_params = self.objects_detector.detect(image)
+ if not (image != self.image).any():
+ continue
+
+ objects_params = objects_detector.detect(image)
with self.condition:
self.image, self.objects_params = image, objects_params
diff --git a/src/polystar/target_pipeline/objects_detectors/trt_model_object_detector.py b/src/polystar/target_pipeline/objects_detectors/trt_model_object_detector.py
index 2ab2154..8ab92f2 100644
--- a/src/polystar/target_pipeline/objects_detectors/trt_model_object_detector.py
+++ b/src/polystar/target_pipeline/objects_detectors/trt_model_object_detector.py
@@ -11,8 +11,6 @@ import tensorrt as trt
from polystar.constants import RESOURCES_DIR
from polystar.models.image import Image
-from polystar.target_pipeline.detected_objects.detected_armor import DetectedArmor
-from polystar.target_pipeline.detected_objects.detected_robot import DetectedRobot
from polystar.target_pipeline.detected_objects.objects_params import ObjectParams
from polystar.target_pipeline.objects_detectors.objects_detector_abc import ObjectsDetectorABC
@@ -24,34 +22,32 @@ class TRTModelObjectsDetector(ObjectsDetectorABC):
def __post_init__(self, model_dir: Path):
self.trt_model = TRTModel(model_dir / "trt_model.bin", (300, 300))
- def detect(self, image: Image) -> Tuple[List[DetectedRobot], List[DetectedArmor]]:
+ def detect(self, image: Image) -> List[ObjectParams]:
results = self.trt_model(image)
- return self._construct_objects_from_trt_results(results, image)
-
- def _construct_objects_from_trt_results(
- self, results: np.ndarray, image: Image
- ) -> Tuple[List[DetectedRobot], List[DetectedArmor]]:
- return self.objects_factory.make_lists(
- [
- ObjectParams(
- ymin=float(ymin),
- xmin=float(xmin),
- ymax=float(ymax),
- xmax=float(xmax),
- score=float(score),
- object_class_id=int(object_class_id),
- )
- for (_, object_class_id, score, xmin, ymin, xmax, ymax) in results
- if object_class_id >= 0
- ],
- image,
+ return _construct_objects_from_trt_results(results)
+
+
+def _construct_objects_from_trt_results(results: np.ndarray) -> List[ObjectParams]:
+ return [
+ ObjectParams(
+ ymin=float(ymin),
+ xmin=float(xmin),
+ ymax=float(ymax),
+ xmax=float(xmax),
+ score=float(score),
+ object_class_id=int(object_class_id),
)
+ for (_, object_class_id, score, xmin, ymin, xmax, ymax) in results
+ if object_class_id == 4 and score >= 0.1
+ ]
class TRTModel:
def __init__(self, trt_model_path: Path, input_size: Tuple[int, int]):
self.input_size = input_size
+ self.cuda_ctx = cuda.Device(0).make_context()
+
self.trt_logger = trt.Logger(trt.Logger.INFO)
self._load_plugins()
self.engine = self._load_engine(trt_model_path)
@@ -117,3 +113,4 @@ class TRTModel:
del self.stream
del self.cuda_outputs
del self.cuda_inputs
+ self.cuda_ctx.pop()
diff --git a/src/research/scripts/demo_pipeline_camera.py b/src/research/scripts/demo_pipeline_camera.py
index cddc032..9a18f56 100644
--- a/src/research/scripts/demo_pipeline_camera.py
+++ b/src/research/scripts/demo_pipeline_camera.py
@@ -28,6 +28,7 @@ class CameraPipelineDemo:
for debug_info in self.pipeline.flow_debug(self.webcam):
self._send_target(debug_info.target)
self._display(viewer, debug_info)
+ self.fps.skip()
def _send_target(self, target: Optional[SimpleTarget]):
if target is not None:
--
GitLab