administrator/frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-05-05 04:57:42 +03:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Fix yolov9 NMS

Browse Source
This commit is contained in:
Nicolas Mowen 2025-04-18 07:19:36 -06:00
parent 19aaa64fe9
commit 6b2246cf08
1 changed files with 11 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

17
frigate/util/model.py
View File

@ -200,9 +200,14 @@ def __post_process_nms_yolo(predictions: np.ndarray, width, height) -> np.ndarra
# Rescale box # Rescale box
boxes = predictions[:, :4] boxes = predictions[:, :4]
boxes_xyxy = np.ones_like(boxes)
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
boxes = boxes_xyxy
input_shape = np.array([width, height, width, height]) # run NMS
boxes = np.divide(boxes, input_shape, dtype=np.float32)
indices = cv2.dnn.NMSBoxes(boxes, scores, score_threshold=0.4, nms_threshold=0.4) indices = cv2.dnn.NMSBoxes(boxes, scores, score_threshold=0.4, nms_threshold=0.4)
detections = np.zeros((20, 6), np.float32) detections = np.zeros((20, 6), np.float32)
for i, (bbox, confidence, class_id) in enumerate( for i, (bbox, confidence, class_id) in enumerate(
@ -214,10 +219,10 @@ def __post_process_nms_yolo(predictions: np.ndarray, width, height) -> np.ndarra
detections[i] = [ detections[i] = [
class_id, class_id,
confidence, confidence,
bbox[1] - bbox[3] / 2, bbox[1] / height,
bbox[0] - bbox[2] / 2, bbox[0] / width,
bbox[1] + bbox[3] / 2, bbox[3] / height,
bbox[0] + bbox[2] / 2, bbox[2] / width,
] ]
return detections return detections