Adding Models

frigate/detectors/plugins/hailo8l.py

@@ -1,6 +1,7 @@
 import logging
 import os
 import urllib.request
+from typing import Optional
 
 import numpy as np
 
@@ -31,38 +32,46 @@ logger = logging.getLogger(__name__)
 # Define the detector key for Hailo
 DETECTOR_KEY = "hailo8l"
 
-
 # Configuration class for model settings
 class ModelConfig(BaseModel):
-    path: str = Field(default=None, title="Model Path")  # Path to the HEF file
+    path: Optional[str] = Field(default=None, title="Model Path")
-
+    type: str = Field(default="ssd_mobilenet_v1", title="Model Type")
+    url: str = Field(default="", title="Model URL")
+    width: int = Field(default=300, title="Model Width")
+    height: int = Field(default=300, title="Model Height")
+    score_threshold: float = Field(default=0.3, title="Score Threshold")
+    max_detections: int = Field(default=30, title="Maximum Detections")
+    input_tensor: str = Field(default="input_tensor", title="Input Tensor Name")
+    input_pixel_format: str = Field(default="RGB", title="Input Pixel Format")
 
 # Configuration class for Hailo detector
 class HailoDetectorConfig(BaseDetectorConfig):
-    type: Literal[DETECTOR_KEY]  # Type of the detector
+    type: Literal[DETECTOR_KEY]
-    device: str = Field(default="PCIe", title="Device Type")  # Device type (e.g., PCIe)
+    device: str = Field(default="PCIe", title="Device Type")
-
+    model: ModelConfig
 
 # Hailo detector class implementation
 class HailoDetector(DetectionApi):
-    type_key = DETECTOR_KEY  # Set the type key to the Hailo detector key
+    type_key = DETECTOR_KEY
 
     def __init__(self, detector_config: HailoDetectorConfig):
-        # Initialize device type and model path from the configuration
+        # Initialize base configuration
         self.h8l_device_type = detector_config.device
         self.h8l_model_path = detector_config.model.path
         self.h8l_model_height = detector_config.model.height
         self.h8l_model_width = detector_config.model.width
-        self.h8l_model_type = detector_config.model.model_type
+        self.h8l_model_type = detector_config.model.type
         self.h8l_tensor_format = detector_config.model.input_tensor
         self.h8l_pixel_format = detector_config.model.input_pixel_format
-        self.model_url = "https://hailo-model-zoo.s3.eu-west-2.amazonaws.com/ModelZoo/Compiled/v2.11.0/hailo8l/ssd_mobilenet_v1.hef"
+        self.model_url = detector_config.model.url
+        self.score_threshold = detector_config.model.score_threshold
+        self.max_detections = detector_config.model.max_detections
+        
         self.cache_dir = "/config/model_cache/h8l_cache"
-        self.expected_model_filename = "ssd_mobilenet_v1.hef"
-        output_type = "FLOAT32"
 
         logger.info(f"Initializing Hailo device as {self.h8l_device_type}")
         self.check_and_prepare_model()
+
         try:
             # Validate device type
             if self.h8l_device_type not in ["PCIe", "M.2"]:
@@ -87,7 +96,8 @@ class HailoDetector(DetectionApi):
                 format_type=self.hef.get_input_vstream_infos()[0].format.type,
             )
             self.output_vstream_params = OutputVStreamParams.make(
-                self.network_group, format_type=getattr(FormatType, output_type)
+                self.network_group,
+                format_type=FormatType.FLOAT32
             )
 
             # Get input and output stream information from the HEF
@@ -99,9 +109,8 @@ class HailoDetector(DetectionApi):
             logger.debug(f"[__init__] Input Tensor Format: {self.h8l_tensor_format}")
             logger.debug(f"[__init__] Input Pixel Format: {self.h8l_pixel_format}")
             logger.debug(f"[__init__] Input VStream Info: {self.input_vstream_info[0]}")
-            logger.debug(
+            logger.debug(f"[__init__] Output VStream Info: {self.output_vstream_info[0]}")
-                f"[__init__] Output VStream Info: {self.output_vstream_info[0]}"
+            
-            )
         except HailoRTException as e:
             logger.error(f"HailoRTException during initialization: {e}")
             raise
@@ -110,12 +119,14 @@ class HailoDetector(DetectionApi):
             raise
 
     def check_and_prepare_model(self):
-        # Ensure cache directory exists
+        """Download and prepare the model if necessary"""
         if not os.path.exists(self.cache_dir):
             os.makedirs(self.cache_dir)
 
-        # Check for the expected model file
+        model_filename = f"{self.h8l_model_type}.hef"
-        model_file_path = os.path.join(self.cache_dir, self.expected_model_filename)
+        model_file_path = os.path.join(self.cache_dir, model_filename)
+        self.h8l_model_path = model_file_path
+
         if not os.path.isfile(model_file_path):
             logger.info(
                 f"A model file was not found at {model_file_path}, Downloading one from {self.model_url}."
@@ -134,9 +145,7 @@ class HailoDetector(DetectionApi):
         )
 
         if tensor_input is None:
-            raise ValueError(
+            raise ValueError("[detect_raw] The 'tensor_input' argument must be provided")
-                "[detect_raw] The 'tensor_input' argument must be provided"
-            )
 
         # Ensure tensor_input is a numpy array
         if isinstance(tensor_input, list):
@@ -182,104 +191,132 @@ class HailoDetector(DetectionApi):
                         logger.error(
                             f"[detect_raw] Missing output stream {self.output_vstream_info[0].name} in inference results"
                         )
-                        return np.zeros((20, 6), np.float32)
+                        return np.zeros((self.max_detections, 6), np.float32)
 
                     raw_output = raw_output[self.output_vstream_info[0].name][0]
                     logger.debug(
                         f"[detect_raw] Raw output for stream {self.output_vstream_info[0].name}: {raw_output}"
                     )
 
-            # Process the raw output
+            # Process the raw output based on model type
             detections = self.process_detections(raw_output)
             if len(detections) == 0:
                 logger.debug(
                     "[detect_raw] No detections found after processing. Setting default values."
                 )
-                return np.zeros((20, 6), np.float32)
+                return np.zeros((self.max_detections, 6), np.float32)
             else:
-                formatted_detections = detections
+                return detections
-                if (
+
-                    formatted_detections.shape[1] != 6
-                ):  # Ensure the formatted detections have 6 columns
-                    logger.error(
-                        f"[detect_raw] Unexpected shape for formatted detections: {formatted_detections.shape}. Expected (20, 6)."
-                    )
-                    return np.zeros((20, 6), np.float32)
-                return formatted_detections
         except HailoRTException as e:
             logger.error(f"[detect_raw] HailoRTException during inference: {e}")
-            return np.zeros((20, 6), np.float32)
+            return np.zeros((self.max_detections, 6), np.float32)
         except Exception as e:
             logger.error(f"[detect_raw] Exception during inference: {e}")
-            return np.zeros((20, 6), np.float32)
+            return np.zeros((self.max_detections, 6), np.float32)
         finally:
             logger.debug("[detect_raw] Exiting function")
 
-    def process_detections(self, raw_detections, threshold=0.5):
+    def process_detections(self, raw_detections, threshold=None):
+        """Process detections based on model type"""
+        if threshold is None:
+            threshold = self.score_threshold
+
+        if self.h8l_model_type == "ssd_mobilenet_v1":
+            return self._process_ssd_detections(raw_detections, threshold)
+        elif self.h8l_model_type == "yolov8s":
+            return self._process_yolo_detections(raw_detections, threshold, version=8)
+        elif self.h8l_model_type == "yolov6n":
+            return self._process_yolo_detections(raw_detections, threshold, version=6)
+        else:
+            logger.error(f"Unsupported model type: {self.h8l_model_type}")
+            return np.zeros((self.max_detections, 6), np.float32)
+
+    def _process_ssd_detections(self, raw_detections, threshold):
+        """Process SSD MobileNet detections"""
         boxes, scores, classes = [], [], []
         num_detections = 0
 
-        logger.debug(f"[process_detections] Raw detections: {raw_detections}")
+        try:
-
+            for detection_set in raw_detections:
-        for i, detection_set in enumerate(raw_detections):
+                if not isinstance(detection_set, np.ndarray) or detection_set.size == 0:
-            if not isinstance(detection_set, np.ndarray) or detection_set.size == 0:
-                logger.debug(
-                    f"[process_detections] Detection set {i} is empty or not an array, skipping."
-                )
-                continue
-
-            logger.debug(
-                f"[process_detections] Detection set {i} shape: {detection_set.shape}"
-            )
-
-            for detection in detection_set:
-                if detection.shape[0] == 0:
-                    logger.debug(
-                        f"[process_detections] Detection in set {i} is empty, skipping."
-                    )
                     continue
 
-                ymin, xmin, ymax, xmax = detection[:4]
+                for detection in detection_set:
-                score = np.clip(detection[4], 0, 1)  # Use np.clip for clarity
+                    if detection.shape[0] == 0:
+                        continue
 
-                if score < threshold:
+                    ymin, xmin, ymax, xmax = detection[:4]
-                    logger.debug(
+                    score = np.clip(detection[4], 0, 1)
-                        f"[process_detections] Detection in set {i} has a score {score} below threshold {threshold}. Skipping."
+
-                    )
+                    if score < threshold:
-                    continue
+                        continue
+
+                    boxes.append([ymin, xmin, ymax, xmax])
+                    scores.append(score)
+                    classes.append(int(detection[5]))
+                    num_detections += 1
+
+            return self._format_output(boxes, scores, classes)
+
+        except Exception as e:
+            logger.error(f"Error processing SSD detections: {e}")
+            return np.zeros((self.max_detections, 6), np.float32)
+
+    def _process_yolo_detections(self, raw_detections, threshold, version):
+        """Process YOLO detections (v6 and v8)"""
+        boxes, scores, classes = [], [], []
+        
+        try:
+            detections = raw_detections[0]
+            
+            for detection in detections:
+                if version == 8:
+                    confidence = detection[4]
+                    if confidence < threshold:
+                        continue
+                    class_scores = detection[5:]
+                else:  # YOLOv6
+                    class_scores = detection[4:]
+                    confidence = np.max(class_scores)
+                    if confidence < threshold:
+                        continue
+
+                x, y, w, h = detection[:4]
+                
+                # Convert to corner format
+                ymin = y - h/2
+                xmin = x - w/2
+                ymax = y + h/2
+                xmax = x + w/2
+                
+                class_id = np.argmax(class_scores)
                 
-                logger.debug(
-                    f"[process_detections] Adding detection with coordinates: ({xmin}, {ymin}), ({xmax}, {ymax}) and score: {score}"
-                )
                 boxes.append([ymin, xmin, ymax, xmax])
-                scores.append(score)
+                scores.append(confidence)
-                classes.append(i)
+                classes.append(class_id)
-                num_detections += 1
 
-        logger.debug(
+            return self._format_output(boxes, scores, classes)
-            f"[process_detections] Boxes: {boxes}, Scores: {scores}, Classes: {classes}, Num detections: {num_detections}"
-        )
 
-        if num_detections == 0:
+        except Exception as e:
-            logger.debug("[process_detections] No valid detections found.")
+            logger.error(f"Error processing YOLO detections: {e}")
-            return np.zeros((20, 6), np.float32)
+            return np.zeros((self.max_detections, 6), np.float32)
 
-        combined = np.hstack(
+    def _format_output(self, boxes, scores, classes):
-            (
+        """Format detections to standard output format"""
-                np.array(classes)[:, np.newaxis],
+        if not boxes:
-                np.array(scores)[:, np.newaxis],
+            return np.zeros((self.max_detections, 6), np.float32)
-                np.array(boxes),
-            )
-        )
             
-        if combined.shape[0] < 20:
+        combined = np.hstack((
-            padding = np.zeros(
+            np.array(classes)[:, np.newaxis],
-                (20 - combined.shape[0], combined.shape[1]), dtype=combined.dtype
+            np.array(scores)[:, np.newaxis],
-            )
+            np.array(boxes)
+        ))
+        
+        if combined.shape[0] < self.max_detections:
+            padding = np.zeros((self.max_detections - combined.shape[0], 6), dtype=np.float32)
             combined = np.vstack((combined, padding))
+        else:
+            combined = combined[:self.max_detections]
             
-        logger.debug(
+        return combined
-            f"[process_detections] Combined detections (padded to 20 if necessary): {np.array_str(combined, precision=4, suppress_small=True)}"
-        )
-
-        return combined[:20, :6]