From e310f3312bc89d6f60be59c654f891431ea10fb5 Mon Sep 17 00:00:00 2001 From: Denis Medyantsev Date: Thu, 9 Mar 2023 17:10:26 +0300 Subject: [PATCH 1/2] Fix numpy data types in model postprocessing. --- savant/deepstream/buffer_processor.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/savant/deepstream/buffer_processor.py b/savant/deepstream/buffer_processor.py index 41fc48895..60a33a66d 100644 --- a/savant/deepstream/buffer_processor.py +++ b/savant/deepstream/buffer_processor.py @@ -451,7 +451,12 @@ def prepare_element_output(self, element: PipelineElement, buffer: Gst.Buffer): # add 0 angle bbox_tensor = np.concatenate( - [bbox_tensor, np.zeros((bbox_tensor.shape[0], 1))], + [ + bbox_tensor, + np.zeros( + (bbox_tensor.shape[0], 1), dtype=np.float32 + ), + ], axis=1, ) @@ -459,7 +464,9 @@ def prepare_element_output(self, element: PipelineElement, buffer: Gst.Buffer): bbox_tensor = np.concatenate( [ bbox_tensor, - np.arange(bbox_tensor.shape[0]).reshape(-1, 1), + np.arange( + bbox_tensor.shape[0], dtype=np.float32 + ).reshape(-1, 1), ], axis=1, ) From 349f91e06434ae6020a8138cc0da4e607cb48b08 Mon Sep 17 00:00:00 2001 From: Denis Medyantsev Date: Thu, 9 Mar 2023 17:10:46 +0300 Subject: [PATCH 2/2] Add numba. --- requirements/base.txt | 5 ++ requirements/deepstream-l4t.txt | 2 - requirements/deepstream.txt | 2 - savant/selector/detector.py | 112 +++++++++++++++++++++----------- 4 files changed, 78 insertions(+), 43 deletions(-) diff --git a/requirements/base.txt b/requirements/base.txt index 1ec6fc270..0e4b66f2f 100644 --- a/requirements/base.txt +++ b/requirements/base.txt @@ -1,3 +1,8 @@ +numpy~=1.22.4 +numba~=0.56.4 + +pycairo~=1.21.0 + fastavro~=1.4.4 pyzmq~=22.2.1 diff --git a/requirements/deepstream-l4t.txt b/requirements/deepstream-l4t.txt index 26c765545..39f7363c1 100644 --- a/requirements/deepstream-l4t.txt +++ b/requirements/deepstream-l4t.txt @@ -1,4 +1,2 @@ -numpy~=1.22.4 -pycairo~=1.21.0 # pyds https://github.com/NVIDIA-AI-IOT/deepstream_python_apps/releases/download/v1.1.6/pyds-1.1.6-py3-none-linux_aarch64.whl \ No newline at end of file diff --git a/requirements/deepstream.txt b/requirements/deepstream.txt index b192c693b..78e7c3aeb 100644 --- a/requirements/deepstream.txt +++ b/requirements/deepstream.txt @@ -1,4 +1,2 @@ -numpy~=1.22.4 -pycairo~=1.21.0 # pyds https://github.com/NVIDIA-AI-IOT/deepstream_python_apps/releases/download/v1.1.6/pyds-1.1.6-py3-none-linux_x86_64.whl diff --git a/savant/selector/detector.py b/savant/selector/detector.py index ef2198ac6..ac9a573ac 100644 --- a/savant/selector/detector.py +++ b/savant/selector/detector.py @@ -1,8 +1,10 @@ """Detector's bbox selectors.""" +from numba import njit, uint8, uint16, float32 import numpy as np from savant.base.selector import BaseSelector +@njit(uint8[:](float32[:, :], float32[:], float32), nogil=True) def nms_cpu( bboxes: np.ndarray, confidences: np.ndarray, threshold: float ) -> np.ndarray: @@ -39,12 +41,71 @@ def nms_cpu( return keep +@njit( + float32[:, :](float32[:, :], float32, float32, uint16, uint16, uint16, uint16), + nogil=True, +) +def default_selector( + bbox_tensor: np.ndarray, + confidence_threshold: float = 0.0, + nms_iou_threshold: float = 0.0, + min_width: int = 0, + min_height: int = 0, + max_width: int = 0, + max_height: int = 0, +) -> np.ndarray: + """Filters bboxes by confidence and size, applies NMS. + + :param bbox_tensor: tensor(class_id, confidence, left, top, width, height) + :param confidence_threshold: confidence threshold + :param nms_iou_threshold: nms iou threshold + :param min_width: minimal bbox width + :param min_height: minimal bbox height + :param max_width: maximum bbox width + :param max_height: maximum bbox height + :return: Selected bbox tensor + """ + selected_bbox_tensor = bbox_tensor.copy() + + if confidence_threshold: + selected_bbox_tensor = selected_bbox_tensor[ + selected_bbox_tensor[:, 1] > confidence_threshold + ] + + if nms_iou_threshold: + keep = nms_cpu( + selected_bbox_tensor[:, 2:6], + selected_bbox_tensor[:, 1], + nms_iou_threshold, + ) + selected_bbox_tensor = selected_bbox_tensor[keep == 1] + + if min_width: + selected_bbox_tensor = selected_bbox_tensor[ + selected_bbox_tensor[:, 4] > min_width + ] + + if min_height: + selected_bbox_tensor = selected_bbox_tensor[ + selected_bbox_tensor[:, 5] > min_height + ] + + if max_width: + selected_bbox_tensor = selected_bbox_tensor[ + selected_bbox_tensor[:, 4] < max_width + ] + + if max_height: + selected_bbox_tensor = selected_bbox_tensor[ + selected_bbox_tensor[:, 5] < max_height + ] + + return selected_bbox_tensor + + class BBoxSelector(BaseSelector): """Detector bbox per class selector. - .. todo:: - - support topk - :param confidence_threshold: confidence threshold :param nms_iou_threshold: nms iou threshold :param min_width: minimal bbox width @@ -77,39 +138,12 @@ def __call__(self, bbox_tensor: np.ndarray) -> np.ndarray: :param bbox_tensor: tensor(class_id, confidence, left, top, width, height) :return: Selected bbox tensor """ - selected_bbox_tensor = bbox_tensor.copy() - - if self.confidence_threshold: - selected_bbox_tensor = selected_bbox_tensor[ - selected_bbox_tensor[:, 1] > self.confidence_threshold - ] - - if self.nms_iou_threshold: - keep = nms_cpu( - selected_bbox_tensor[:, 2:6], - selected_bbox_tensor[:, 1], - self.nms_iou_threshold, - ) - selected_bbox_tensor = selected_bbox_tensor[keep == 1] - - if self.min_width: - selected_bbox_tensor = selected_bbox_tensor[ - selected_bbox_tensor[:, 4] > self.min_width - ] - - if self.min_height: - selected_bbox_tensor = selected_bbox_tensor[ - selected_bbox_tensor[:, 5] > self.min_height - ] - - if self.max_width: - selected_bbox_tensor = selected_bbox_tensor[ - selected_bbox_tensor[:, 4] < self.max_width - ] - - if self.max_height: - selected_bbox_tensor = selected_bbox_tensor[ - selected_bbox_tensor[:, 5] < self.max_height - ] - - return selected_bbox_tensor + return default_selector( + bbox_tensor=bbox_tensor, + confidence_threshold=self.confidence_threshold, + nms_iou_threshold=self.nms_iou_threshold, + min_width=self.min_width, + min_height=self.min_height, + max_width=self.max_width, + max_height=self.max_height, + )