From 3d9b1996f9b274124879af294f4ac371d34f5118 Mon Sep 17 00:00:00 2001 From: Sergey Krashevich Date: Wed, 26 Apr 2023 02:26:54 +0300 Subject: [PATCH] Add MetaDetector implementation and configuration schema --- frigate/detectors/plugins/meta.py | 108 ++++++++++++++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 frigate/detectors/plugins/meta.py diff --git a/frigate/detectors/plugins/meta.py b/frigate/detectors/plugins/meta.py new file mode 100644 index 000000000..c2bb4dca6 --- /dev/null +++ b/frigate/detectors/plugins/meta.py @@ -0,0 +1,108 @@ +import logging +import numpy as np +import os +from frigate.detectors.detection_api import DetectionApi +from frigate.detectors.detector_config import BaseDetectorConfig, ModelConfig +from frigate.util import deep_merge +from typing import List, Tuple, Dict, Any, Literal +from typing import Union +from typing_extensions import Annotated +from enum import Enum +from pydantic import Field, parse_obj_as +import importlib +import pkgutil + +logger = logging.getLogger(__name__) + +DETECTOR_KEY = "meta_detector" + +DetectorConfig = Annotated[ + Union[tuple(BaseDetectorConfig.__subclasses__())], + Field(discriminator="type"), + ] + +class MetaDetectorConfig(BaseDetectorConfig): + type: Literal[DETECTOR_KEY] + detectors: Dict[str, DetectorConfig] = Field( + default={"cpu": {"type": "cpu"}}, + title="Detector hardware configuration.", + ) + max_detections: int = Field( + default=20, + title="Maximum number of detections to return after merging results", + ) + + +class MetaDetector(DetectionApi): + type_key = DETECTOR_KEY + + def __init__(self, meta_detector_config: MetaDetectorConfig): + self.max_detections = meta_detector_config.max_detections + + self.detectors = [] + + for key, detector in meta_detector_config.detectors.items(): + detector_config: DetectorConfig = parse_obj_as(DetectorConfig, detector) + if detector_config.model is None: + detector_config.model = meta_detector_config.model + else: + model = detector_config.model + schema = ModelConfig.schema()["properties"] + if ( + model.width != schema["width"]["default"] + or model.height != schema["height"]["default"] + or model.labelmap_path is not None + or model.labelmap is not {} + or model.input_tensor != schema["input_tensor"]["default"] + or model.input_pixel_format + != schema["input_pixel_format"]["default"] + ): + logger.warning( + "Customizing more than a detector model path is unsupported." + ) + merged_model = deep_merge( + detector_config.model.dict(exclude_unset=True), + meta_detector_config.model.dict(exclude_unset=True), + ) + detector_config.model = ModelConfig.parse_obj(merged_model) + meta_detector_config.detectors[key] = detector_config + self.detectors.append(self.create_detector(detector_config)) + + + + def merge_detections(self, detections_list: List[np.ndarray]) -> np.ndarray: + all_detections = np.vstack(detections_list) + sorted_indices = np.argsort(-all_detections[:, 1]) + return all_detections[sorted_indices[: self.max_detections]] + + def detect_raw(self, tensor_input) -> np.ndarray: + detections_list = [ + detector.detect_raw(tensor_input) for detector in self.detectors + ] + return self.merge_detections(detections_list) + + def create_detector(self, detector_config): + + current_module_name = os.path.splitext(os.path.basename(__file__))[0] + modules_folder = os.path.dirname(os.path.abspath(__file__)) + module_prefix = __package__ + "." + + _included_modules = [module for module in pkgutil.iter_modules([modules_folder], module_prefix) if module.name != current_module_name] + + plugin_modules = [] + + for _, name, _ in _included_modules: + try: + # currently openvino may fail when importing + # on an arm device with 64 KiB page size. + plugin_modules.append(importlib.import_module(name)) + except ImportError as e: + logger.error(f"Error importing detector runtime: {e}") + + + + api_types = {det.type_key: det for det in DetectionApi.__subclasses__()} + api = api_types.get(detector_config.type) + if not api: + raise ValueError(detector_config.type) + return api(detector_config)