From 7693b55dc76737bf48700a5f999426b03c71bceb Mon Sep 17 00:00:00 2001 From: HonzaCuhel Date: Fri, 30 Aug 2024 16:59:19 +0200 Subject: [PATCH] Fix arg typing --- depthai_nodes/ml/parsers/utils/yolo.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/depthai_nodes/ml/parsers/utils/yolo.py b/depthai_nodes/ml/parsers/utils/yolo.py index 56011a1..14e6d73 100644 --- a/depthai_nodes/ml/parsers/utils/yolo.py +++ b/depthai_nodes/ml/parsers/utils/yolo.py @@ -51,7 +51,7 @@ def non_max_suppression( max_nms: int = 30000, max_wh: int = 7680, kpts_mode: bool = False, -) -> list[np.ndarray]: +) -> List[np.ndarray]: """Performs Non-Maximum Suppression (NMS) on inference results. @param prediction: Prediction from the model, shape = (batch_size, boxes, xy+wh+...) @@ -79,7 +79,7 @@ def non_max_suppression( @param kpts_mode: Keypoints mode. @type kpts_mode: bool @return: An array of detections with either kpts or segmentation outputs. - @rtype: list[np.ndarray] + @rtype: List[np.ndarray] """ bs = prediction.shape[0] # batch size # Keypoints: 4 (bbox) + 1 (objectness) + 51 (kpts) = 56