diff --git a/frigate/detectors/detector_config.py b/frigate/detectors/detector_config.py index be12e7fcc..ce7738493 100644 --- a/frigate/detectors/detector_config.py +++ b/frigate/detectors/detector_config.py @@ -33,11 +33,12 @@ class InputDTypeEnum(str, Enum): class ModelTypeEnum(str, Enum): + dfine = "dfine" + rfdetr = "rfdetr" ssd = "ssd" yolox = "yolox" yolov9 = "yolov9" yolonas = "yolonas" - dfine = "dfine" yologeneric = "yolo-generic" diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index d94b4660f..03940f418 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/frigate/detectors/plugins/onnx.py @@ -11,6 +11,7 @@ from frigate.detectors.detector_config import ( ) from frigate.util.model import ( get_ort_providers, + post_process_rfdetr, post_process_dfine, post_process_yolov9, ) @@ -73,7 +74,9 @@ class ONNXDetector(DetectionApi): model_input_name = self.model.get_inputs()[0].name tensor_output = self.model.run(None, {model_input_name: tensor_input}) - if self.onnx_model_type == ModelTypeEnum.yolonas: + if self.onnx_model_type == ModelTypeEnum.rfdetr: + return post_process_rfdetr(tensor_output) + elif self.onnx_model_type == ModelTypeEnum.yolonas: predictions = tensor_output[0] detections = np.zeros((20, 6), np.float32) diff --git a/frigate/util/model.py b/frigate/util/model.py index d96493ee6..55a5de86c 100644 --- a/frigate/util/model.py +++ b/frigate/util/model.py @@ -13,7 +13,11 @@ logger = logging.getLogger(__name__) ### Post Processing -def post_process_dfine(tensor_output: np.ndarray, width, height) -> np.ndarray: + + +def post_process_dfine( + tensor_output: np.ndarray, width: int, height: int +) -> np.ndarray: class_ids = tensor_output[0][tensor_output[2] > 0.4] boxes = tensor_output[1][tensor_output[2] > 0.4] scores = tensor_output[2][tensor_output[2] > 0.4] @@ -41,6 +45,61 @@ def post_process_dfine(tensor_output: np.ndarray, width, height) -> np.ndarray: return detections +def post_process_rfdetr(tensor_output: list[np.ndarray, np.ndarray]) -> np.ndarray: + boxes = tensor_output[0] + raw_scores = tensor_output[1] + + # apply soft max to scores + exp = np.exp(raw_scores - np.max(raw_scores, axis=-1, keepdims=True)) + all_scores = exp / np.sum(exp, axis=-1, keepdims=True) + + # get highest scoring class from every detection + scores = np.max(all_scores[0, :, 1:], axis=-1) + labels = np.argmax(all_scores[0, :, 1:], axis=-1) + + idxs = scores > 0.4 + filtered_boxes = boxes[0, idxs] + filtered_scores = scores[idxs] + filtered_labels = labels[idxs] + + # convert boxes from [x_center, y_center, width, height] + x_center, y_center, w, h = ( + filtered_boxes[:, 0], + filtered_boxes[:, 1], + filtered_boxes[:, 2], + filtered_boxes[:, 3], + ) + x_min = x_center - w / 2 + y_min = y_center - h / 2 + x_max = x_center + w / 2 + y_max = y_center + h / 2 + filtered_boxes = np.stack([x_min, y_min, x_max, y_max], axis=-1) + + # apply nms + indices = cv2.dnn.NMSBoxes( + filtered_boxes, filtered_scores, score_threshold=0.4, nms_threshold=0.4 + ) + detections = np.zeros((20, 6), np.float32) + + for i, (bbox, confidence, class_id) in enumerate( + zip(filtered_boxes[indices], filtered_scores[indices], filtered_labels[indices]) + ): + if i == 20: + break + + detections[i] = [ + class_id, + confidence, + bbox[1], + bbox[0], + bbox[3], + bbox[2], + ] + # print(f"found a detection {detections[i]}") + + return detections + + def post_process_yolov9(predictions: np.ndarray, width, height) -> np.ndarray: predictions = np.squeeze(predictions).T scores = np.max(predictions[:, 4:], axis=1)