diff --git a/frigate/detectors/detection_api.py b/frigate/detectors/detection_api.py index ccccc1b65d..31dbd3671b 100644 --- a/frigate/detectors/detection_api.py +++ b/frigate/detectors/detection_api.py @@ -16,7 +16,7 @@ class DetectionApi(ABC): @abstractmethod def __init__(self, detector_config: BaseDetectorConfig): self.detector_config = detector_config - self.thresh = 0.5 + self.thresh = 0.4 self.height = detector_config.model.height self.width = detector_config.model.width @@ -24,58 +24,21 @@ class DetectionApi(ABC): def detect_raw(self, tensor_input): pass - def post_process_yolonas(self, output): - """ - @param output: output of inference - expected shape: [np.array(1, N, 4), np.array(1, N, 80)] - where N depends on the input size e.g. N=2100 for 320x320 images + def calculate_grids_strides(self) -> None: + grids = [] + expanded_strides = [] - @return: best results: np.array(20, 6) where each row is - in this order (class_id, score, y1/height, x1/width, y2/height, x2/width) - """ + # decode and orient predictions + strides = [8, 16, 32] + hsizes = [self.height // stride for stride in strides] + wsizes = [self.width // stride for stride in strides] - N = output[0].shape[1] + 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)) - boxes = output[0].reshape(N, 4) - scores = output[1].reshape(N, 80) - - class_ids = np.argmax(scores, axis=1) - scores = scores[np.arange(N), class_ids] - - args_best = np.argwhere(scores > self.thresh)[:, 0] - - num_matches = len(args_best) - if num_matches == 0: - return np.zeros((20, 6), np.float32) - elif num_matches > 20: - args_best20 = np.argpartition(scores[args_best], -20)[-20:] - args_best = args_best[args_best20] - - boxes = boxes[args_best] - class_ids = class_ids[args_best] - scores = scores[args_best] - - boxes = np.transpose( - np.vstack( - ( - boxes[:, 1] / self.height, - boxes[:, 0] / self.width, - boxes[:, 3] / self.height, - boxes[:, 2] / self.width, - ) - ) - ) - - results = np.hstack( - (class_ids[..., np.newaxis], scores[..., np.newaxis], boxes) - ) - - return np.resize(results, (20, 6)) - - def post_process(self, output): - if self.detector_config.model.model_type == ModelTypeEnum.yolonas: - return self.post_process_yolonas(output) - else: - raise ValueError( - f'Model type "{self.detector_config.model.model_type}" is currently not supported.' - ) + self.grids = np.concatenate(grids, 1) + self.expanded_strides = np.concatenate(expanded_strides, 1) diff --git a/frigate/detectors/plugins/onnx.py b/frigate/detectors/plugins/onnx.py index 2f9b3d8a46..6c6a39e783 100644 --- a/frigate/detectors/plugins/onnx.py +++ b/frigate/detectors/plugins/onnx.py @@ -60,23 +60,7 @@ class ONNXDetector(DetectionApi): 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) + self.calculate_grids_strides() logger.info(f"ONNX: {path} loaded") diff --git a/frigate/detectors/plugins/openvino.py b/frigate/detectors/plugins/openvino.py index 9c7ed5248d..b90a22a7b4 100644 --- a/frigate/detectors/plugins/openvino.py +++ b/frigate/detectors/plugins/openvino.py @@ -133,25 +133,7 @@ class OvDetector(DetectionApi): break self.num_classes = tensor_shape[2] - 5 logger.info(f"YOLOX model has {self.num_classes} classes") - self.set_strides_grids() - - def set_strides_grids(self): - grids = [] - expanded_strides = [] - - strides = [8, 16, 32] - - hsize_list = [self.h // stride for stride in strides] - wsize_list = [self.w // stride for stride in strides] - - for hsize, wsize, stride in zip(hsize_list, wsize_list, 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) + self.calculate_grids_strides() ## Takes in class ID, confidence score, and array of [x, y, w, h] that describes detection position, ## returns an array that's easily passable back to Frigate.