diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index aef6e909b7..2f9b3d8a46 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/frigate/detectors/plugins/onnx.py @@ -14,6 +14,7 @@ from frigate.util.model import ( post_process_dfine, post_process_rfdetr, post_process_yolo, + post_process_yolox, ) logger = logging.getLogger(__name__) @@ -58,6 +59,25 @@ class ONNXDetector(DetectionApi): self.onnx_model_shape = detector_config.model.input_tensor path = detector_config.model.path + if self.onnx_model_type == ModelTypeEnum.yolox: + grids = [] + expanded_strides = [] + + # decode and orient predictions + strides = [8, 16, 32] + hsizes = [self.h // stride for stride in strides] + wsizes = [self.w // stride for stride in strides] + + for hsize, wsize, stride in zip(hsizes, wsizes, strides): + xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) + grid = np.stack((xv, yv), 2).reshape(1, -1, 2) + grids.append(grid) + shape = grid.shape[:2] + expanded_strides.append(np.full((*shape, 1), stride)) + + self.grids = np.concatenate(grids, 1) + self.expanded_strides = np.concatenate(expanded_strides, 1) + logger.info(f"ONNX: {path} loaded") def detect_raw(self, tensor_input: np.ndarray): @@ -99,6 +119,10 @@ class ONNXDetector(DetectionApi): return detections elif self.onnx_model_type == ModelTypeEnum.yologeneric: return post_process_yolo(tensor_output, self.w, self.h) + elif self.onnx_model_type == ModelTypeEnum.yolox: + return post_process_yolox( + tensor_output[0], self.w, self.h, self.grids, self.expanded_strides + ) else: raise Exception( f"{self.onnx_model_type} is currently not supported for onnx. See the docs for more info on supported models." diff --git a/frigate/util/model.py b/frigate/util/model.py index b459e3f740..4888575c6d 100644 --- a/frigate/util/model.py +++ b/frigate/util/model.py @@ -230,24 +230,13 @@ def post_process_yolo(output: list[np.ndarray], width: int, height: int) -> np.n return __post_process_nms_yolo(output[0], width, height) -def post_process_yolox(predictions: np.ndarray, width: int, height: int) -> np.ndarray: - grids = [] - expanded_strides = [] - - # decode and orient predictions - strides = [8, 16, 32] - hsizes = [height // stride for stride in strides] - wsizes = [width // stride for stride in strides] - - for hsize, wsize, stride in zip(hsizes, wsizes, strides): - xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) - grid = np.stack((xv, yv), 2).reshape(1, -1, 2) - grids.append(grid) - shape = grid.shape[:2] - expanded_strides.append(np.full((*shape, 1), stride)) - - grids = np.concatenate(grids, 1) - expanded_strides = np.concatenate(expanded_strides, 1) +def post_process_yolox( + predictions: np.ndarray, + width: int, + height: int, + grids: np.ndarray, + expanded_strides: np.ndarray, +) -> np.ndarray: predictions[..., :2] = (predictions[..., :2] + grids) * expanded_strides predictions[..., 2:4] = np.exp(predictions[..., 2:4]) * expanded_strides @@ -269,15 +258,6 @@ def post_process_yolox(predictions: np.ndarray, width: int, height: int) -> np.n boxes_xyxy, scores, score_threshold=0.4, nms_threshold=0.4 ) - final_boxes = boxes_xyxy[indices] - final_scores = scores[indices] - final_cls_inds = cls_inds[indices] - - print(f"frig boxes: {final_boxes}") - print(f"frig cls: {final_cls_inds}") - print(f"frig scores: {final_scores}") - - detections = np.zeros((20, 6), np.float32) for i, (bbox, confidence, class_id) in enumerate( zip(boxes_xyxy[indices], scores[indices], cls_inds[indices]) @@ -288,10 +268,10 @@ def post_process_yolox(predictions: np.ndarray, width: int, height: int) -> np.n detections[i] = [ class_id, confidence, - bbox[1], - bbox[0], - bbox[3], - bbox[2], + bbox[1] / height, + bbox[0] / width, + bbox[3] / height, + bbox[2] / width, ] print(f"got {detections[i]}")