mirror of
https://github.com/blakeblackshear/frigate.git
synced 2025-03-09 00:16:54 +01:00
394 lines
16 KiB
Python
Executable File
394 lines
16 KiB
Python
Executable File
import logging
|
|
import os
|
|
import subprocess
|
|
import urllib.request
|
|
import numpy as np
|
|
import queue
|
|
import threading
|
|
from functools import partial
|
|
from typing import Dict, Optional, List, Tuple
|
|
|
|
try:
|
|
from hailo_platform import (
|
|
HEF,
|
|
ConfigureParams,
|
|
FormatType,
|
|
HailoRTException,
|
|
HailoStreamInterface,
|
|
InputVStreamParams,
|
|
OutputVStreamParams,
|
|
VDevice,
|
|
HailoSchedulingAlgorithm,
|
|
)
|
|
except ModuleNotFoundError:
|
|
pass
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing_extensions import Literal
|
|
|
|
from frigate.const import MODEL_CACHE_DIR
|
|
from frigate.detectors.detection_api import DetectionApi
|
|
from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum, InputTensorEnum, PixelFormatEnum, InputDTypeEnum
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
# ----------------- Inline Utility Functions ----------------- #
|
|
def preprocess_image(image: Image.Image, model_w: int, model_h: int) -> Image.Image:
|
|
"""
|
|
Resize image with unchanged aspect ratio using padding.
|
|
"""
|
|
img_w, img_h = image.size
|
|
scale = min(model_w / img_w, model_h / img_h)
|
|
new_img_w, new_img_h = int(img_w * scale), int(img_h * scale)
|
|
image = image.resize((new_img_w, new_img_h), Image.Resampling.BICUBIC)
|
|
padded_image = Image.new('RGB', (model_w, model_h), (114, 114, 114))
|
|
padded_image.paste(image, ((model_w - new_img_w) // 2, (model_h - new_img_h) // 2))
|
|
return padded_image
|
|
|
|
def extract_detections(input_data: list, threshold: float = 0.5) -> dict:
|
|
"""
|
|
(Legacy extraction function; not used by detect_raw below.)
|
|
Extract detections from raw inference output.
|
|
"""
|
|
boxes, scores, classes = [], [], []
|
|
num_detections = 0
|
|
for i, detection in enumerate(input_data):
|
|
if len(detection) == 0:
|
|
continue
|
|
for det in detection:
|
|
bbox, score = det[:4], det[4]
|
|
if score >= threshold:
|
|
boxes.append(bbox)
|
|
scores.append(score)
|
|
classes.append(i)
|
|
num_detections += 1
|
|
return {
|
|
'detection_boxes': boxes,
|
|
'detection_classes': classes,
|
|
'detection_scores': scores,
|
|
'num_detections': num_detections
|
|
}
|
|
# ----------------- End of Utility Functions ----------------- #
|
|
|
|
# Global constants and default URLs
|
|
DETECTOR_KEY = "hailo8l"
|
|
ARCH = None
|
|
H8_DEFAULT_MODEL = "yolov8s.hef"
|
|
H8L_DEFAULT_MODEL = "yolov6n.hef"
|
|
H8_DEFAULT_URL = "https://hailo-model-zoo.s3.eu-west-2.amazonaws.com/ModelZoo/Compiled/v2.14.0/hailo8/yolov8s.hef"
|
|
H8L_DEFAULT_URL = "https://hailo-model-zoo.s3.eu-west-2.amazonaws.com/ModelZoo/Compiled/v2.14.0/hailo8l/yolov6n.hef"
|
|
|
|
def detect_hailo_arch():
|
|
try:
|
|
result = subprocess.run(['hailortcli', 'fw-control', 'identify'], capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
print(f"Error running hailortcli: {result.stderr}")
|
|
return None
|
|
for line in result.stdout.split('\n'):
|
|
if "Device Architecture" in line:
|
|
if "HAILO8L" in line:
|
|
return "hailo8l"
|
|
elif "HAILO8" in line:
|
|
return "hailo8"
|
|
print("Could not determine Hailo architecture from device information.")
|
|
return None
|
|
except Exception as e:
|
|
print(f"An error occurred while detecting Hailo architecture: {e}")
|
|
return None
|
|
|
|
# ----------------- Inline Asynchronous Inference Class ----------------- #
|
|
class HailoAsyncInference:
|
|
def __init__(
|
|
self,
|
|
hef_path: str,
|
|
input_queue: queue.Queue,
|
|
output_queue: queue.Queue,
|
|
batch_size: int = 1,
|
|
input_type: Optional[str] = None,
|
|
output_type: Optional[Dict[str, str]] = None,
|
|
send_original_frame: bool = False,
|
|
) -> None:
|
|
self.input_queue = input_queue
|
|
self.output_queue = output_queue
|
|
|
|
# Create VDevice parameters with round-robin scheduling
|
|
params = VDevice.create_params()
|
|
params.scheduling_algorithm = HailoSchedulingAlgorithm.ROUND_ROBIN
|
|
|
|
# Load HEF and create the infer model
|
|
self.hef = HEF(hef_path)
|
|
self.target = VDevice(params)
|
|
self.infer_model = self.target.create_infer_model(hef_path)
|
|
self.infer_model.set_batch_size(batch_size)
|
|
if input_type is not None:
|
|
self._set_input_type(input_type)
|
|
if output_type is not None:
|
|
self._set_output_type(output_type)
|
|
self.output_type = output_type
|
|
self.send_original_frame = send_original_frame
|
|
|
|
def _set_input_type(self, input_type: Optional[str] = None) -> None:
|
|
self.infer_model.input().set_format_type(getattr(FormatType, input_type))
|
|
|
|
def _set_output_type(self, output_type_dict: Optional[Dict[str, str]] = None) -> None:
|
|
for output_name, output_type in output_type_dict.items():
|
|
self.infer_model.output(output_name).set_format_type(getattr(FormatType, output_type))
|
|
|
|
def callback(self, completion_info, bindings_list: List, input_batch: List):
|
|
if completion_info.exception:
|
|
logging.error(f"Inference error: {completion_info.exception}")
|
|
else:
|
|
for i, bindings in enumerate(bindings_list):
|
|
if len(bindings._output_names) == 1:
|
|
result = bindings.output().get_buffer()
|
|
else:
|
|
result = {
|
|
name: np.expand_dims(bindings.output(name).get_buffer(), axis=0)
|
|
for name in bindings._output_names
|
|
}
|
|
self.output_queue.put((input_batch[i], result))
|
|
|
|
def _create_bindings(self, configured_infer_model) -> object:
|
|
if self.output_type is None:
|
|
output_buffers = {
|
|
output_info.name: np.empty(
|
|
self.infer_model.output(output_info.name).shape,
|
|
dtype=getattr(np, str(output_info.format.type).split(".")[1].lower())
|
|
)
|
|
for output_info in self.hef.get_output_vstream_infos()
|
|
}
|
|
else:
|
|
output_buffers = {
|
|
name: np.empty(
|
|
self.infer_model.output(name).shape,
|
|
dtype=getattr(np, self.output_type[name].lower())
|
|
)
|
|
for name in self.output_type
|
|
}
|
|
return configured_infer_model.create_bindings(output_buffers=output_buffers)
|
|
|
|
def get_input_shape(self) -> Tuple[int, ...]:
|
|
return self.hef.get_input_vstream_infos()[0].shape
|
|
|
|
def run(self) -> None:
|
|
# Configure the infer model once and reuse vstream settings via run_async
|
|
with self.infer_model.configure() as configured_infer_model:
|
|
while True:
|
|
batch_data = self.input_queue.get()
|
|
if batch_data is None:
|
|
break # Sentinel to exit loop
|
|
if self.send_original_frame:
|
|
original_batch, preprocessed_batch = batch_data
|
|
else:
|
|
preprocessed_batch = batch_data
|
|
bindings_list = []
|
|
for frame in preprocessed_batch:
|
|
bindings = self._create_bindings(configured_infer_model)
|
|
bindings.input().set_buffer(np.array(frame))
|
|
bindings_list.append(bindings)
|
|
configured_infer_model.wait_for_async_ready(timeout_ms=10000)
|
|
job = configured_infer_model.run_async(
|
|
bindings_list,
|
|
partial(
|
|
self.callback,
|
|
input_batch=original_batch if self.send_original_frame else preprocessed_batch,
|
|
bindings_list=bindings_list,
|
|
)
|
|
)
|
|
job.wait(10000) # Wait for the last job to complete
|
|
# ----------------- End of Async Class ----------------- #
|
|
|
|
# ----------------- HailoDetector Class ----------------- #
|
|
class HailoDetector(DetectionApi):
|
|
type_key = DETECTOR_KEY
|
|
|
|
def __init__(self, detector_config: 'HailoDetectorConfig'):
|
|
global ARCH
|
|
ARCH = detect_hailo_arch()
|
|
self.cache_dir = MODEL_CACHE_DIR
|
|
self.device_type = detector_config.device
|
|
# Model attributes should be provided in detector_config.model
|
|
self.model_path = detector_config.model.path if hasattr(detector_config.model, "path") else None
|
|
self.model_height = detector_config.model.height if hasattr(detector_config.model, "height") else None
|
|
self.model_width = detector_config.model.width if hasattr(detector_config.model, "width") else None
|
|
self.model_type = detector_config.model.model_type if hasattr(detector_config.model, "model_type") else None
|
|
self.tensor_format = detector_config.model.input_tensor if hasattr(detector_config.model, "input_tensor") else None
|
|
self.pixel_format = detector_config.model.input_pixel_format if hasattr(detector_config.model, "input_pixel_format") else None
|
|
self.input_dtype = detector_config.model.input_dtype if hasattr(detector_config.model, "input_dtype") else None
|
|
self.url = detector_config.url
|
|
self.output_type = "FLOAT32"
|
|
self.working_model_path = self.check_and_prepare()
|
|
|
|
# Set up asynchronous inference
|
|
self.batch_size = 1
|
|
self.input_queue = queue.Queue()
|
|
self.output_queue = queue.Queue()
|
|
try:
|
|
logging.debug(f"[INIT] Loading HEF model from {self.working_model_path}")
|
|
self.inference_engine = HailoAsyncInference(
|
|
self.working_model_path,
|
|
self.input_queue,
|
|
self.output_queue,
|
|
self.batch_size
|
|
)
|
|
self.input_shape = self.inference_engine.get_input_shape()
|
|
logging.debug(f"[INIT] Model input shape: {self.input_shape}")
|
|
except Exception as e:
|
|
logging.error(f"[INIT] Failed to initialize HailoAsyncInference: {e}")
|
|
raise
|
|
|
|
@staticmethod
|
|
def extract_model_name(path: str = None, url: str = None) -> str:
|
|
model_name = None
|
|
if path and path.endswith(".hef"):
|
|
model_name = os.path.basename(path)
|
|
elif url and url.endswith(".hef"):
|
|
model_name = os.path.basename(url)
|
|
else:
|
|
print("Model name not found in path or URL. Checking default settings...")
|
|
if ARCH == "hailo8":
|
|
model_name = H8_DEFAULT_MODEL
|
|
else:
|
|
model_name = H8L_DEFAULT_MODEL
|
|
print(f"Using default model: {model_name}")
|
|
return model_name
|
|
|
|
@staticmethod
|
|
def download_model(url: str, destination: str):
|
|
if not url.endswith(".hef"):
|
|
raise ValueError("Invalid model URL. Only .hef files are supported.")
|
|
try:
|
|
urllib.request.urlretrieve(url, destination)
|
|
print(f"Downloaded model to {destination}")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to download model from {url}: {str(e)}")
|
|
|
|
def check_and_prepare(self) -> str:
|
|
if not os.path.exists(self.cache_dir):
|
|
os.makedirs(self.cache_dir)
|
|
model_name = self.extract_model_name(self.model_path, self.url)
|
|
model_path = os.path.join(self.cache_dir, model_name)
|
|
if not self.model_path and not self.url:
|
|
if os.path.exists(model_path):
|
|
print(f"Model found in cache: {model_path}")
|
|
return model_path
|
|
else:
|
|
print(f"Downloading default model: {model_name}")
|
|
if ARCH == "hailo8":
|
|
self.download_model(H8_DEFAULT_URL, model_path)
|
|
else:
|
|
self.download_model(H8L_DEFAULT_URL, model_path)
|
|
elif self.model_path and self.url:
|
|
if os.path.exists(self.model_path):
|
|
print(f"Model found at path: {self.model_path}")
|
|
return self.model_path
|
|
else:
|
|
print(f"Model not found at path. Downloading from URL: {self.url}")
|
|
self.download_model(self.url, model_path)
|
|
elif self.url:
|
|
print(f"Downloading model from URL: {self.url}")
|
|
self.download_model(self.url, model_path)
|
|
elif self.model_path:
|
|
if os.path.exists(self.model_path):
|
|
print(f"Using existing model at: {self.model_path}")
|
|
return self.model_path
|
|
else:
|
|
raise FileNotFoundError(f"Model file not found at: {self.model_path}")
|
|
return model_path
|
|
|
|
def detect_raw(self, tensor_input):
|
|
logging.debug("[DETECT_RAW] Starting detection")
|
|
# Ensure tensor_input has a batch dimension
|
|
if isinstance(tensor_input, np.ndarray) and len(tensor_input.shape) == 3:
|
|
tensor_input = np.expand_dims(tensor_input, axis=0)
|
|
logging.debug(f"[DETECT_RAW] Expanded input shape to {tensor_input.shape}")
|
|
|
|
# Enqueue input and a sentinel value
|
|
self.input_queue.put(tensor_input)
|
|
self.input_queue.put(None) # Sentinel value
|
|
|
|
# Run the inference engine
|
|
self.inference_engine.run()
|
|
result = self.output_queue.get()
|
|
if result is None:
|
|
logging.error("[DETECT_RAW] No inference result received")
|
|
return np.zeros((20, 6), dtype=np.float32)
|
|
|
|
original_input, infer_results = result
|
|
logging.debug("[DETECT_RAW] Inference completed.")
|
|
|
|
# If infer_results is a single-element list, unwrap it.
|
|
if isinstance(infer_results, list) and len(infer_results) == 1:
|
|
infer_results = infer_results[0]
|
|
|
|
# Set your threshold (adjust as needed)
|
|
threshold = 0.4
|
|
all_detections = []
|
|
|
|
# Loop over the output list (each element corresponds to one output stream)
|
|
for idx, detection_set in enumerate(infer_results):
|
|
# Skip empty arrays
|
|
if not isinstance(detection_set, np.ndarray) or detection_set.size == 0:
|
|
continue
|
|
|
|
logging.debug(f"[DETECT_RAW] Processing detection set {idx} with shape {detection_set.shape}")
|
|
# For each detection row in the set:
|
|
for det in detection_set:
|
|
# Expecting at least 5 elements: [ymin, xmin, ymax, xmax, confidence]
|
|
if det.shape[0] < 5:
|
|
continue
|
|
score = float(det[4])
|
|
if score < threshold:
|
|
continue
|
|
# If there is a 6th element, assume it's a class id; otherwise use dummy class 0.
|
|
if det.shape[0] >= 6:
|
|
cls = int(det[5])
|
|
else:
|
|
cls = 0
|
|
# Append in the order Frigate expects: [class_id, confidence, ymin, xmin, ymax, xmax]
|
|
all_detections.append([cls, score, det[0], det[1], det[2], det[3]])
|
|
|
|
# If no valid detections were found, return a zero array.
|
|
if len(all_detections) == 0:
|
|
logging.warning("[DETECT_RAW] No valid detections found.")
|
|
return np.zeros((20, 6), dtype=np.float32)
|
|
|
|
detections_array = np.array(all_detections, dtype=np.float32)
|
|
|
|
# Pad or truncate to exactly 20 rows
|
|
if detections_array.shape[0] < 20:
|
|
pad = np.zeros((20 - detections_array.shape[0], 6), dtype=np.float32)
|
|
detections_array = np.vstack((detections_array, pad))
|
|
elif detections_array.shape[0] > 20:
|
|
detections_array = detections_array[:20, :]
|
|
|
|
logging.debug(f"[DETECT_RAW] Processed detections: {detections_array}")
|
|
return detections_array
|
|
|
|
# Preprocess method using inline utility
|
|
def preprocess(self, image):
|
|
return preprocess_image(image, self.input_shape[1], self.input_shape[0])
|
|
|
|
# Close the Hailo device
|
|
def close(self):
|
|
logging.debug("[CLOSE] Closing HailoDetector")
|
|
try:
|
|
self.inference_engine.hef.close()
|
|
logging.debug("Hailo device closed successfully")
|
|
except Exception as e:
|
|
logging.error(f"Failed to close Hailo device: {e}")
|
|
raise
|
|
|
|
# Asynchronous detection wrapper
|
|
def async_detect(self, tensor_input, callback):
|
|
def detection_thread():
|
|
result = self.detect_raw(tensor_input)
|
|
callback(result)
|
|
thread = threading.Thread(target=detection_thread)
|
|
thread.start()
|
|
|
|
# ----------------- Configuration Class ----------------- #
|
|
class HailoDetectorConfig(BaseDetectorConfig):
|
|
type: Literal[DETECTOR_KEY]
|
|
device: str = Field(default="PCIe", title="Device Type")
|
|
url: Optional[str] = Field(default=None, title="Custom Model URL")
|