mirror of
https://github.com/blakeblackshear/frigate.git
synced 2025-08-04 13:47:37 +02:00
Final Async Update
This commit is contained in:
parent
eadd55eec6
commit
7813476500
@ -1,322 +1,192 @@
|
|||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
|
import queue
|
||||||
|
import threading
|
||||||
|
import subprocess
|
||||||
import urllib.request
|
import urllib.request
|
||||||
from typing import Optional
|
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
from hailo_platform import (
|
||||||
try:
|
HEF,
|
||||||
from hailo_platform import (
|
ConfigureParams,
|
||||||
HEF,
|
FormatType,
|
||||||
ConfigureParams,
|
HailoRTException,
|
||||||
FormatType,
|
HailoStreamInterface,
|
||||||
HailoRTException,
|
VDevice,
|
||||||
HailoStreamInterface,
|
HailoSchedulingAlgorithm,
|
||||||
InferVStreams,
|
InferVStreams,
|
||||||
InputVStreamParams,
|
InputVStreamParams,
|
||||||
OutputVStreamParams,
|
OutputVStreamParams
|
||||||
VDevice,
|
)
|
||||||
)
|
|
||||||
except ModuleNotFoundError:
|
|
||||||
pass
|
|
||||||
|
|
||||||
from pydantic import BaseModel, Field
|
|
||||||
from typing_extensions import Literal
|
|
||||||
|
|
||||||
from frigate.detectors.detection_api import DetectionApi
|
from frigate.detectors.detection_api import DetectionApi
|
||||||
from frigate.detectors.detector_config import BaseDetectorConfig
|
from frigate.detectors.detector_config import BaseDetectorConfig
|
||||||
|
from pydantic import BaseModel, Field
|
||||||
|
from typing_extensions import Literal
|
||||||
|
from typing import Optional
|
||||||
|
|
||||||
# Set up logging
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
# Define the detector key for Hailo
|
|
||||||
DETECTOR_KEY = "hailo8l"
|
DETECTOR_KEY = "hailo8l"
|
||||||
|
|
||||||
# Configuration class for model settings
|
def get_device_architecture():
|
||||||
|
"""Get the device architecture from hailortcli."""
|
||||||
|
try:
|
||||||
|
result = subprocess.run(['hailortcli', 'fw-control', 'identify'], capture_output=True, text=True)
|
||||||
|
for line in result.stdout.split('\n'):
|
||||||
|
if "Device Architecture" in line:
|
||||||
|
return line.split(':')[1].strip().lower()
|
||||||
|
except Exception:
|
||||||
|
return "unknown"
|
||||||
|
|
||||||
class ModelConfig(BaseModel):
|
class ModelConfig(BaseModel):
|
||||||
path: Optional[str] = Field(default=None, title="Model Path")
|
path: Optional[str] = Field(default=None, title="Model Path")
|
||||||
type: str = Field(default="ssd_mobilenet_v1", title="Model Type")
|
type: str = Field(default="yolov8s", title="Model Type")
|
||||||
url: str = Field(default="", title="Model URL")
|
width: int = Field(default=640, title="Model Width")
|
||||||
width: int = Field(default=300, title="Model Width")
|
height: int = Field(default=640, title="Model Height")
|
||||||
height: int = Field(default=300, title="Model Height")
|
|
||||||
score_threshold: float = Field(default=0.3, title="Score Threshold")
|
score_threshold: float = Field(default=0.3, title="Score Threshold")
|
||||||
max_detections: int = Field(default=30, title="Maximum Detections")
|
max_detections: int = Field(default=30, title="Maximum Detections")
|
||||||
input_tensor: str = Field(default="input_tensor", title="Input Tensor Name")
|
input_tensor: str = Field(default="input_tensor", title="Input Tensor Name")
|
||||||
input_pixel_format: str = Field(default="RGB", title="Input Pixel Format")
|
input_pixel_format: str = Field(default="RGB", title="Input Pixel Format")
|
||||||
|
|
||||||
# Configuration class for Hailo detector
|
|
||||||
class HailoDetectorConfig(BaseDetectorConfig):
|
class HailoDetectorConfig(BaseDetectorConfig):
|
||||||
type: Literal[DETECTOR_KEY]
|
type: Literal[DETECTOR_KEY]
|
||||||
device: str = Field(default="PCIe", title="Device Type")
|
device: str = Field(default="PCIe", title="Device Type")
|
||||||
model: ModelConfig
|
model: ModelConfig
|
||||||
|
|
||||||
# Hailo detector class implementation
|
class HailoAsyncInference:
|
||||||
|
def __init__(self, config: HailoDetectorConfig):
|
||||||
|
self.config = config
|
||||||
|
self.input_queue = queue.Queue()
|
||||||
|
self.output_queue = queue.Queue()
|
||||||
|
params = VDevice.create_params()
|
||||||
|
params.scheduling_algorithm = HailoSchedulingAlgorithm.ROUND_ROBIN
|
||||||
|
self.target = VDevice(params)
|
||||||
|
self.hef = HEF(self.config.model.path)
|
||||||
|
self.infer_model = self.target.create_infer_model(self.config.model.path)
|
||||||
|
self.infer_model.set_batch_size(1)
|
||||||
|
|
||||||
|
def infer(self):
|
||||||
|
while True:
|
||||||
|
batch_data = self.input_queue.get()
|
||||||
|
if batch_data is None:
|
||||||
|
break
|
||||||
|
|
||||||
|
bindings = []
|
||||||
|
for frame in batch_data:
|
||||||
|
binding = self.infer_model.create_bindings()
|
||||||
|
binding.input().set_buffer(frame)
|
||||||
|
bindings.append(binding)
|
||||||
|
|
||||||
|
self.infer_model.run_async(bindings, self._callback, batch_data)
|
||||||
|
|
||||||
|
def _callback(self, completion_info, bindings_list, batch_data):
|
||||||
|
if completion_info.exception:
|
||||||
|
logger.error(f"Inference error: {completion_info.exception}")
|
||||||
|
else:
|
||||||
|
results = [binding.output().get_buffer() for binding in bindings_list]
|
||||||
|
self.output_queue.put((batch_data, results))
|
||||||
|
|
||||||
|
def stop(self):
|
||||||
|
self.input_queue.put(None)
|
||||||
|
|
||||||
class HailoDetector(DetectionApi):
|
class HailoDetector(DetectionApi):
|
||||||
type_key = DETECTOR_KEY
|
type_key = DETECTOR_KEY
|
||||||
|
|
||||||
def __init__(self, detector_config: HailoDetectorConfig):
|
def __init__(self, config: HailoDetectorConfig):
|
||||||
# Initialize base configuration
|
super().__init__()
|
||||||
self.h8l_device_type = detector_config.device
|
self.async_inference = HailoAsyncInference(config)
|
||||||
self.h8l_model_path = detector_config.model.path
|
self.worker_thread = threading.Thread(target=self.async_inference.infer)
|
||||||
self.h8l_model_height = detector_config.model.height
|
self.worker_thread.start()
|
||||||
self.h8l_model_width = detector_config.model.width
|
|
||||||
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 = 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"
|
|
||||||
|
|
||||||
logger.info(f"Initializing Hailo device as {self.h8l_device_type}")
|
# Determine device architecture
|
||||||
|
self.device_architecture = get_device_architecture()
|
||||||
|
if self.device_architecture not in ["hailo8", "hailo8l"]:
|
||||||
|
raise RuntimeError(f"Unsupported device architecture: {self.device_architecture}")
|
||||||
|
logger.info(f"Device architecture detected: {self.device_architecture}")
|
||||||
|
|
||||||
|
# Ensure the model is available
|
||||||
|
self.cache_dir = "/config/model_cache/h8l_cache"
|
||||||
|
self.expected_model_filename = f"{config.model.type}.hef"
|
||||||
self.check_and_prepare_model()
|
self.check_and_prepare_model()
|
||||||
|
|
||||||
try:
|
|
||||||
# Validate device type
|
|
||||||
if self.h8l_device_type not in ["PCIe", "M.2"]:
|
|
||||||
raise ValueError(f"Unsupported device type: {self.h8l_device_type}")
|
|
||||||
|
|
||||||
# Initialize the Hailo device
|
|
||||||
self.target = VDevice()
|
|
||||||
# Load the HEF (Hailo's binary format for neural networks)
|
|
||||||
self.hef = HEF(self.h8l_model_path)
|
|
||||||
# Create configuration parameters from the HEF
|
|
||||||
self.configure_params = ConfigureParams.create_from_hef(
|
|
||||||
hef=self.hef, interface=HailoStreamInterface.PCIe
|
|
||||||
)
|
|
||||||
# Configure the device with the HEF
|
|
||||||
self.network_groups = self.target.configure(self.hef, self.configure_params)
|
|
||||||
self.network_group = self.network_groups[0]
|
|
||||||
self.network_group_params = self.network_group.create_params()
|
|
||||||
|
|
||||||
# Create input and output virtual stream parameters
|
|
||||||
self.input_vstream_params = InputVStreamParams.make(
|
|
||||||
self.network_group,
|
|
||||||
format_type=self.hef.get_input_vstream_infos()[0].format.type,
|
|
||||||
)
|
|
||||||
self.output_vstream_params = OutputVStreamParams.make(
|
|
||||||
self.network_group,
|
|
||||||
format_type=FormatType.FLOAT32
|
|
||||||
)
|
|
||||||
|
|
||||||
# Get input and output stream information from the HEF
|
|
||||||
self.input_vstream_info = self.hef.get_input_vstream_infos()
|
|
||||||
self.output_vstream_info = self.hef.get_output_vstream_infos()
|
|
||||||
|
|
||||||
logger.info("Hailo device initialized successfully")
|
|
||||||
logger.debug(f"[__init__] Model Path: {self.h8l_model_path}")
|
|
||||||
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(f"[__init__] Output VStream Info: {self.output_vstream_info[0]}")
|
|
||||||
|
|
||||||
except HailoRTException as e:
|
|
||||||
logger.error(f"HailoRTException during initialization: {e}")
|
|
||||||
raise
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Failed to initialize Hailo device: {e}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
def check_and_prepare_model(self):
|
def check_and_prepare_model(self):
|
||||||
"""Download and prepare the model if necessary"""
|
# Ensure cache directory exists
|
||||||
if not os.path.exists(self.cache_dir):
|
if not os.path.exists(self.cache_dir):
|
||||||
os.makedirs(self.cache_dir)
|
os.makedirs(self.cache_dir)
|
||||||
|
|
||||||
model_filename = f"{self.h8l_model_type}.hef"
|
# Check for the expected model file
|
||||||
model_file_path = os.path.join(self.cache_dir, model_filename)
|
model_file_path = os.path.join(self.cache_dir, self.expected_model_filename)
|
||||||
self.h8l_model_path = model_file_path
|
self.async_inference.config.model.path = model_file_path
|
||||||
|
|
||||||
if not os.path.isfile(model_file_path):
|
if not os.path.isfile(model_file_path):
|
||||||
logger.info(
|
if self.async_inference.config.model.path:
|
||||||
f"A model file was not found at {model_file_path}, Downloading one from {self.model_url}."
|
logger.info(
|
||||||
)
|
f"A model file was not found at {model_file_path}, Downloading one from the provided URL."
|
||||||
urllib.request.urlretrieve(self.model_url, model_file_path)
|
)
|
||||||
logger.info(f"A model file was downloaded to {model_file_path}.")
|
urllib.request.urlretrieve(self.async_inference.config.model.path, model_file_path)
|
||||||
|
logger.info(f"A model file was downloaded to {model_file_path}.")
|
||||||
|
else:
|
||||||
|
raise RuntimeError("Model file path is missing and no URL is provided.")
|
||||||
else:
|
else:
|
||||||
logger.info(
|
logger.info(
|
||||||
f"A model file already exists at {model_file_path} not downloading one."
|
f"A model file already exists at {model_file_path} not downloading one."
|
||||||
)
|
)
|
||||||
|
|
||||||
def detect_raw(self, tensor_input):
|
def detect_raw(self, tensor_input):
|
||||||
logger.debug("[detect_raw] Entering function")
|
"""
|
||||||
logger.debug(
|
Perform inference and return raw detection results.
|
||||||
f"[detect_raw] The `tensor_input` = {tensor_input} tensor_input shape = {tensor_input.shape}"
|
"""
|
||||||
)
|
preprocessed_input = self.preprocess(tensor_input)
|
||||||
|
self.async_inference.input_queue.put([preprocessed_input])
|
||||||
if tensor_input is None:
|
|
||||||
raise ValueError("[detect_raw] The 'tensor_input' argument must be provided")
|
|
||||||
|
|
||||||
# Ensure tensor_input is a numpy array
|
|
||||||
if isinstance(tensor_input, list):
|
|
||||||
tensor_input = np.array(tensor_input)
|
|
||||||
logger.debug(
|
|
||||||
f"[detect_raw] Converted tensor_input to numpy array: shape {tensor_input.shape}"
|
|
||||||
)
|
|
||||||
|
|
||||||
input_data = tensor_input
|
|
||||||
logger.debug(
|
|
||||||
f"[detect_raw] Input data for inference shape: {tensor_input.shape}, dtype: {tensor_input.dtype}"
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
with InferVStreams(
|
batch_data, raw_results = self.async_inference.output_queue.get(timeout=5)
|
||||||
self.network_group,
|
return self.postprocess(raw_results)
|
||||||
self.input_vstream_params,
|
except queue.Empty:
|
||||||
self.output_vstream_params,
|
logger.warning("Inference timed out")
|
||||||
) as infer_pipeline:
|
return np.zeros((20, 6), np.float32)
|
||||||
input_dict = {}
|
|
||||||
if isinstance(input_data, dict):
|
|
||||||
input_dict = input_data
|
|
||||||
logger.debug("[detect_raw] it a dictionary.")
|
|
||||||
elif isinstance(input_data, (list, tuple)):
|
|
||||||
for idx, layer_info in enumerate(self.input_vstream_info):
|
|
||||||
input_dict[layer_info.name] = input_data[idx]
|
|
||||||
logger.debug("[detect_raw] converted from list/tuple.")
|
|
||||||
else:
|
|
||||||
if len(input_data.shape) == 3:
|
|
||||||
input_data = np.expand_dims(input_data, axis=0)
|
|
||||||
logger.debug("[detect_raw] converted from an array.")
|
|
||||||
input_dict[self.input_vstream_info[0].name] = input_data
|
|
||||||
|
|
||||||
logger.debug(
|
def preprocess(self, frame):
|
||||||
f"[detect_raw] Input dictionary for inference keys: {input_dict.keys()}"
|
input_shape = (self.async_inference.hef.get_input_vstream_infos()[0].shape)
|
||||||
)
|
resized_frame = np.resize(frame, input_shape)
|
||||||
|
return resized_frame / 255.0
|
||||||
|
|
||||||
with self.network_group.activate(self.network_group_params):
|
def postprocess(self, raw_output):
|
||||||
raw_output = infer_pipeline.infer(input_dict)
|
model_type = self.async_inference.config.model.type
|
||||||
logger.debug(f"[detect_raw] Raw inference output: {raw_output}")
|
if model_type == "ssd_mobilenet_v1":
|
||||||
|
return self._process_ssd(raw_output)
|
||||||
if self.output_vstream_info[0].name not in raw_output:
|
elif model_type in ["yolov8s", "yolov8m", "yolov6n"]:
|
||||||
logger.error(
|
return self._process_yolo(raw_output, version=model_type[-1])
|
||||||
f"[detect_raw] Missing output stream {self.output_vstream_info[0].name} in inference results"
|
|
||||||
)
|
|
||||||
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 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((self.max_detections, 6), np.float32)
|
|
||||||
else:
|
|
||||||
return detections
|
|
||||||
|
|
||||||
except HailoRTException as e:
|
|
||||||
logger.error(f"[detect_raw] HailoRTException during inference: {e}")
|
|
||||||
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((self.max_detections, 6), np.float32)
|
|
||||||
finally:
|
|
||||||
logger.debug("[detect_raw] Exiting function")
|
|
||||||
|
|
||||||
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:
|
else:
|
||||||
logger.error(f"Unsupported model type: {self.h8l_model_type}")
|
logger.error(f"Unsupported model type: {model_type}")
|
||||||
return np.zeros((self.max_detections, 6), np.float32)
|
return []
|
||||||
|
|
||||||
def _process_ssd_detections(self, raw_detections, threshold):
|
def _process_ssd(self, raw_output):
|
||||||
"""Process SSD MobileNet detections"""
|
detections = []
|
||||||
boxes, scores, classes = [], [], []
|
for detection in raw_output[1]:
|
||||||
num_detections = 0
|
score = detection[4]
|
||||||
|
if score >= self.async_inference.config.model.score_threshold:
|
||||||
try:
|
ymin, xmin, ymax, xmax = detection[:4]
|
||||||
for detection_set in raw_detections:
|
detections.append({
|
||||||
if not isinstance(detection_set, np.ndarray) or detection_set.size == 0:
|
"bounding_box": [xmin, ymin, xmax, ymax],
|
||||||
continue
|
"score": score,
|
||||||
|
"class": int(detection[5])
|
||||||
for detection in detection_set:
|
})
|
||||||
if detection.shape[0] == 0:
|
return detections
|
||||||
continue
|
|
||||||
|
|
||||||
ymin, xmin, ymax, xmax = detection[:4]
|
|
||||||
score = np.clip(detection[4], 0, 1)
|
|
||||||
|
|
||||||
if score < threshold:
|
|
||||||
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
|
|
||||||
|
|
||||||
|
def _process_yolo(self, raw_output, version):
|
||||||
|
detections = []
|
||||||
|
for detection in raw_output[1]:
|
||||||
|
confidence = detection[4] if version == "8" else np.max(detection[5:])
|
||||||
|
if confidence >= self.async_inference.config.model.score_threshold:
|
||||||
x, y, w, h = detection[:4]
|
x, y, w, h = detection[:4]
|
||||||
|
ymin, xmin, ymax, xmax = y - h / 2, x - w / 2, y + h / 2, x + w / 2
|
||||||
# Convert to corner format
|
class_id = np.argmax(detection[5:])
|
||||||
ymin = y - h/2
|
detections.append({
|
||||||
xmin = x - w/2
|
"bounding_box": [xmin, ymin, xmax, ymax],
|
||||||
ymax = y + h/2
|
"score": confidence,
|
||||||
xmax = x + w/2
|
"class": class_id
|
||||||
|
})
|
||||||
class_id = np.argmax(class_scores)
|
return detections
|
||||||
|
|
||||||
boxes.append([ymin, xmin, ymax, xmax])
|
|
||||||
scores.append(confidence)
|
|
||||||
classes.append(class_id)
|
|
||||||
|
|
||||||
return self._format_output(boxes, scores, classes)
|
def stop(self):
|
||||||
|
self.async_inference.stop()
|
||||||
except Exception as e:
|
self.worker_thread.join()
|
||||||
logger.error(f"Error processing YOLO detections: {e}")
|
|
||||||
return np.zeros((self.max_detections, 6), np.float32)
|
|
||||||
|
|
||||||
def _format_output(self, boxes, scores, classes):
|
|
||||||
"""Format detections to standard output format"""
|
|
||||||
if not boxes:
|
|
||||||
return np.zeros((self.max_detections, 6), np.float32)
|
|
||||||
|
|
||||||
combined = np.hstack((
|
|
||||||
np.array(classes)[:, np.newaxis],
|
|
||||||
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]
|
|
||||||
|
|
||||||
return combined
|
|
||||||
|
Loading…
Reference in New Issue
Block a user