From 515597345aef8bb15e75887d3ff5afecacca73ba Mon Sep 17 00:00:00 2001 From: Tim Date: Wed, 13 Aug 2025 17:55:45 -0400 Subject: [PATCH] ran ruff --- frigate/detectors/plugins/memryx.py | 107 ++++++++++++++++------------ frigate/object_detection/base.py | 14 ++-- 2 files changed, 65 insertions(+), 56 deletions(-) diff --git a/frigate/detectors/plugins/memryx.py b/frigate/detectors/plugins/memryx.py index 391561632..0a948c880 100644 --- a/frigate/detectors/plugins/memryx.py +++ b/frigate/detectors/plugins/memryx.py @@ -18,13 +18,18 @@ except ModuleNotFoundError: from pydantic import BaseModel, Field from typing_extensions import Literal from frigate.detectors.detection_api import DetectionApi -from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum, InputTensorEnum +from frigate.detectors.detector_config import ( + BaseDetectorConfig, + ModelTypeEnum, + InputTensorEnum, +) from frigate.util.model import post_process_yolo logger = logging.getLogger(__name__) DETECTOR_KEY = "memryx" + # Configuration class for model settings class ModelConfig(BaseModel): path: str = Field(default=None, title="Model Path") # Path to the DFP file @@ -37,7 +42,7 @@ class MemryXDetectorConfig(BaseDetectorConfig): class MemryXDetector(DetectionApi): - type_key = DETECTOR_KEY # Set the type key + type_key = DETECTOR_KEY # Set the type key supported_models = [ ModelTypeEnum.ssd, ModelTypeEnum.yolonas, @@ -54,7 +59,9 @@ class MemryXDetector(DetectionApi): if "model_type" in getattr(model_cfg, "__fields_set__", set()): detector_config.model.model_type = model_cfg.model_type else: - logger.info("model_type not set in config — defaulting to yolonas for MemryX.") + logger.info( + "model_type not set in config — defaulting to yolonas for MemryX." + ) detector_config.model.model_type = ModelTypeEnum.yolonas self.capture_queue = Queue(maxsize=10) @@ -65,13 +72,13 @@ class MemryXDetector(DetectionApi): self.memx_model_path = detector_config.model.path # Path to .dfp file self.memx_post_model = None # Path to .post file self.expected_post_model = None - + self.memx_device_path = detector_config.device # Device path # Parse the device string to split PCIe: device_str = self.memx_device_path self.device_id = [] self.device_id.append(int(device_str.split(":")[1])) - + self.memx_model_height = detector_config.model.height self.memx_model_width = detector_config.model.width self.memx_model_type = detector_config.model.model_type @@ -80,41 +87,51 @@ class MemryXDetector(DetectionApi): if self.memx_model_type == ModelTypeEnum.yologeneric: model_mapping = { - (640, 640): ("https://developer.memryx.com/example_files/2p0_frigate/yolov9_640.zip", "yolov9_640"), - (320, 320): ("https://developer.memryx.com/example_files/2p0_frigate/yolov9_320.zip", "yolov9_320") + (640, 640): ( + "https://developer.memryx.com/example_files/2p0_frigate/yolov9_640.zip", + "yolov9_640", + ), + (320, 320): ( + "https://developer.memryx.com/example_files/2p0_frigate/yolov9_320.zip", + "yolov9_320", + ), } self.model_url, self.model_folder = model_mapping.get( (self.memx_model_height, self.memx_model_width), - ("https://developer.memryx.com/example_files/2p0_frigate/yolov9_320.zip", "yolov9_320") - ) - self.expected_dfp_model = ( - "YOLO_v9_small_onnx.dfp" + ( + "https://developer.memryx.com/example_files/2p0_frigate/yolov9_320.zip", + "yolov9_320", + ), ) + self.expected_dfp_model = "YOLO_v9_small_onnx.dfp" elif self.memx_model_type == ModelTypeEnum.yolonas: model_mapping = { - (640, 640): ("https://developer.memryx.com/example_files/2p0_frigate/yolonas_640.zip", "yolonas_640"), - (320, 320): ("https://developer.memryx.com/example_files/2p0_frigate/yolonas_320.zip", "yolonas_320") + (640, 640): ( + "https://developer.memryx.com/example_files/2p0_frigate/yolonas_640.zip", + "yolonas_640", + ), + (320, 320): ( + "https://developer.memryx.com/example_files/2p0_frigate/yolonas_320.zip", + "yolonas_320", + ), } self.model_url, self.model_folder = model_mapping.get( (self.memx_model_height, self.memx_model_width), - ("https://developer.memryx.com/example_files/2p0_frigate/yolonas_320.zip", "yolonas_320") - ) - self.expected_dfp_model = ( - "yolo_nas_s.dfp" - ) - self.expected_post_model = ( - "yolo_nas_s_post.onnx" + ( + "https://developer.memryx.com/example_files/2p0_frigate/yolonas_320.zip", + "yolonas_320", + ), ) + self.expected_dfp_model = "yolo_nas_s.dfp" + self.expected_post_model = "yolo_nas_s_post.onnx" elif self.memx_model_type == ModelTypeEnum.yolox: self.model_folder = "yolox" self.model_url = ( "https://developer.memryx.com/example_files/2p0_frigate/yolox.zip" ) - self.expected_dfp_model = ( - "YOLOX_640_640_3_onnx.dfp" - ) + self.expected_dfp_model = "YOLOX_640_640_3_onnx.dfp" self.set_strides_grids() elif self.memx_model_type == ModelTypeEnum.ssd: @@ -122,12 +139,8 @@ class MemryXDetector(DetectionApi): self.model_url = ( "https://developer.memryx.com/example_files/2p0_frigate/ssd.zip" ) - self.expected_dfp_model = ( - "SSDlite_MobileNet_v2_320_320_3_onnx.dfp" - ) - self.expected_post_model = ( - "SSDlite_MobileNet_v2_320_320_3_onnx_post.onnx" - ) + self.expected_dfp_model = "SSDlite_MobileNet_v2_320_320_3_onnx.dfp" + self.expected_post_model = "SSDlite_MobileNet_v2_320_320_3_onnx_post.onnx" self.check_and_prepare_model() logger.info( @@ -143,9 +156,9 @@ class MemryXDetector(DetectionApi): self.accl = AsyncAccl( self.memx_model_path, device_ids=self.device_id, # AsyncAccl device ids - local_mode=True + local_mode=True, ) - + # Models that use cropped post-processing sections (YOLO-NAS and SSD) # --> These will be moved to pure numpy in the future to improve performance on low-end CPUs if self.memx_post_model: @@ -161,19 +174,13 @@ class MemryXDetector(DetectionApi): except Exception as e: logger.error(f"Failed to initialize MemryX model: {e}") raise - + def load_yolo_constants(self): base = f"{self.cache_dir}/{self.model_folder}" # constants for yolov9 post-processing - self.const_A = np.load( - f"{base}/_model_22_Constant_9_output_0.npy" - ) - self.const_B = np.load( - f"{base}/_model_22_Constant_10_output_0.npy" - ) - self.const_C = np.load( - f"{base}/_model_22_Constant_12_output_0.npy" - ) + self.const_A = np.load(f"{base}/_model_22_Constant_9_output_0.npy") + self.const_B = np.load(f"{base}/_model_22_Constant_10_output_0.npy") + self.const_C = np.load(f"{base}/_model_22_Constant_12_output_0.npy") def check_and_prepare_model(self): """Check if models exist; if not, download and extract them.""" @@ -182,7 +189,11 @@ class MemryXDetector(DetectionApi): model_subdir = os.path.join(self.cache_dir, self.model_folder) dfp_path = os.path.join(model_subdir, self.expected_dfp_model) - post_path = os.path.join(model_subdir, self.expected_post_model) if self.expected_post_model else None + post_path = ( + os.path.join(model_subdir, self.expected_post_model) + if self.expected_post_model + else None + ) dfp_exists = os.path.exists(dfp_path) post_exists = os.path.exists(post_path) if post_path else True @@ -210,7 +221,11 @@ class MemryXDetector(DetectionApi): # Re-assign model paths after extraction self.memx_model_path = os.path.join(model_subdir, self.expected_dfp_model) - self.memx_post_model = os.path.join(model_subdir, self.expected_post_model) if self.expected_post_model else None + self.memx_post_model = ( + os.path.join(model_subdir, self.expected_post_model) + if self.expected_post_model + else None + ) if self.memx_model_type == ModelTypeEnum.yologeneric: self.load_yolo_constants() @@ -232,7 +247,9 @@ class MemryXDetector(DetectionApi): if self.memx_model_type == ModelTypeEnum.yolonas: if tensor_input.ndim == 4 and tensor_input.shape[1:] == (320, 320, 3): logger.debug("Transposing tensor from NHWC to NCHW for YOLO-NAS") - tensor_input = np.transpose(tensor_input, (0, 3, 1, 2)) # (1, H, W, C) → (1, C, H, W) + tensor_input = np.transpose( + tensor_input, (0, 3, 1, 2) + ) # (1, H, W, C) → (1, C, H, W) tensor_input = tensor_input.astype(np.float32) tensor_input /= 255 @@ -390,7 +407,6 @@ class MemryXDetector(DetectionApi): return reshaped def post_process_yolox(self, output): - output_785 = output[0] # 785 output_794 = output[1] # 794 output_795 = output[2] # 795 @@ -528,7 +544,6 @@ class MemryXDetector(DetectionApi): def process_output(self, *outputs): """Output callback function -- receives frames from the MX3 and triggers post-processing""" if self.memx_model_type == ModelTypeEnum.yologeneric: - conv_out1 = outputs[0] conv_out2 = outputs[1] conv_out3 = outputs[2] diff --git a/frigate/object_detection/base.py b/frigate/object_detection/base.py index 0d308ab9f..2f08e9175 100644 --- a/frigate/object_detection/base.py +++ b/frigate/object_detection/base.py @@ -97,10 +97,10 @@ class AsyncLocalObjectDetector(BaseLocalDetector): def async_send_input(self, tensor_input: np.ndarray, connection_id): tensor_input = self._transform_input(tensor_input) return self.detect_api.send_input(connection_id, tensor_input) - + def async_receive_output(self): return self.detect_api.receive_output() - + def prepare_detector(name, out_events): threading.current_thread().name = f"detector:{name}" @@ -136,10 +136,7 @@ def run_detector( start: Value, detector_config: BaseDetectorConfig, ): - - stop_event, frame_manager, outputs, logger = prepare_detector( - name, out_events - ) + stop_event, frame_manager, outputs, logger = prepare_detector(name, out_events) object_detector = LocalObjectDetector(detector_config=detector_config) @@ -179,10 +176,7 @@ def async_run_detector( start: Value, detector_config: BaseDetectorConfig, ): - - stop_event, frame_manager, outputs, logger = prepare_detector( - name, out_events - ) + stop_event, frame_manager, outputs, logger = prepare_detector(name, out_events) object_detector = AsyncLocalObjectDetector(detector_config=detector_config)