mirror of
https://github.com/blakeblackshear/frigate.git
synced 2026-08-02 00:52:17 +03:00
Auto convert ONNX models to RKNN format (#19674)
* Implement base rknn conversion * Remove unused * Formatting * Add model conversion lock so it doesn't break when multiple detectors are defined * Ignore unused impor t
This commit is contained in:
@@ -12,6 +12,7 @@ from frigate.const import MODEL_CACHE_DIR
|
||||
from frigate.detectors.detection_api import DetectionApi
|
||||
from frigate.detectors.detector_config import BaseDetectorConfig, ModelTypeEnum
|
||||
from frigate.util.model import post_process_yolo
|
||||
from frigate.util.rknn_converter import auto_convert_model
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -94,7 +95,31 @@ class Rknn(DetectionApi):
|
||||
# user provided models should be a path and contain a "/"
|
||||
if "/" in model_path:
|
||||
model_props["preset"] = False
|
||||
model_props["path"] = model_path
|
||||
|
||||
# Check if this is an ONNX model or model without extension that needs conversion
|
||||
if model_path.endswith(".onnx") or not os.path.splitext(model_path)[1]:
|
||||
# Try to auto-convert to RKNN format
|
||||
logger.info(
|
||||
f"Attempting to auto-convert {model_path} to RKNN format..."
|
||||
)
|
||||
|
||||
# Determine model type from config
|
||||
model_type = self.detector_config.model.model_type
|
||||
|
||||
# Auto-convert the model
|
||||
converted_path = auto_convert_model(model_path, model_type.value)
|
||||
|
||||
if converted_path:
|
||||
model_props["path"] = converted_path
|
||||
logger.info(f"Successfully converted model to: {converted_path}")
|
||||
else:
|
||||
# Fall back to original path if conversion fails
|
||||
logger.warning(
|
||||
f"Failed to convert {model_path} to RKNN format, using original path"
|
||||
)
|
||||
model_props["path"] = model_path
|
||||
else:
|
||||
model_props["path"] = model_path
|
||||
else:
|
||||
model_props["preset"] = True
|
||||
|
||||
|
||||
Reference in New Issue
Block a user