Mirrors/blakeblackshear.frigate
1
0
Fork 0
mirror of https://github.com/blakeblackshear/frigate.git synced 2026-04-28 23:06:13 +02:00
Code Issues Projects Releases Wiki Activity

Add automatic RKNN conversion and support for semantic search model (#19676)

Browse Source 
* Create RKNN model runner and and use for jina v1 clip

* Formatting

* Handle model type inference

* Properly provide input to RKNN

* Adjust rknn conversion

* Update docs

* Formatting

* Fix path handling

* Handle inputs

* Cleanup

* Change normalization for better accuracy

* Clarify supported models

* Remove testing
This commit is contained in:
Nicolas Mowen
2025-08-21 05:30:14 -06:00
committed by GitHub
parent efeb089ff8
commit 1be84d6833
4 changed files with 233 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

83
frigate/util/rknn_converter.py
View File

@@ -27,9 +27,50 @@ MODEL_TYPE_CONFIGS = {
"std_values": [[255, 255, 255]],
"target_platform": None, # Will be set dynamically
},
"jina-clip-v1-vision": {
"mean_values": [[0.48145466 * 255, 0.4578275 * 255, 0.40821073 * 255]],
"std_values": [[0.26862954 * 255, 0.26130258 * 255, 0.27577711 * 255]],
"target_platform": None, # Will be set dynamically
},
}
def get_rknn_model_type(model_path: str) -> str | None:
if all(keyword in str(model_path) for keyword in ["jina-clip-v1", "vision"]):
return "jina-clip-v1-vision"
model_name = os.path.basename(str(model_path)).lower()
if any(keyword in model_name for keyword in ["yolo", "yolox", "yolonas"]):
return model_name
return None
def is_rknn_compatible(model_path: str, model_type: str | None = None) -> bool:
"""
Check if a model is compatible with RKNN conversion.
Args:
model_path: Path to the model file
model_type: Type of the model (if known)
Returns:
True if the model is RKNN-compatible, False otherwise
"""
soc = get_soc_type()
if soc is None:
return False
if not model_type:
model_type = get_rknn_model_type(model_path)
if model_type and model_type in MODEL_TYPE_CONFIGS:
return True
return False
def ensure_torch_dependencies() -> bool:
"""Dynamically install torch dependencies if not available."""
try:
@@ -67,13 +108,12 @@ def ensure_torch_dependencies() -> bool:
def ensure_rknn_toolkit() -> bool:
"""Ensure RKNN toolkit is available."""
try:
import rknn # type: ignore # noqa: F401
from rknn.api import RKNN # type: ignore # noqa: F401
logger.debug("RKNN toolkit is already available")
return True
except ImportError:
logger.error("RKNN toolkit not found. Please ensure it's installed.")
except ImportError as e:
logger.error(f"RKNN toolkit not found. Please ensure it's installed. {e}")
return False
@@ -109,11 +149,11 @@ def convert_onnx_to_rknn(
True if conversion successful, False otherwise
"""
if not ensure_torch_dependencies():
logger.error("PyTorch dependencies not available")
logger.debug("PyTorch dependencies not available")
return False
if not ensure_rknn_toolkit():
logger.error("RKNN toolkit not available")
logger.debug("RKNN toolkit not available")
return False
# Get SoC type if not provided
@@ -125,7 +165,7 @@ def convert_onnx_to_rknn(
# Get model config for the specified type
if model_type not in MODEL_TYPE_CONFIGS:
logger.error(f"Unsupported model type: {model_type}")
logger.debug(f"Unsupported model type: {model_type}")
return False
config = MODEL_TYPE_CONFIGS[model_type].copy()
@@ -138,7 +178,16 @@ def convert_onnx_to_rknn(
rknn = RKNN(verbose=True)
rknn.config(**config)
if rknn.load_onnx(model=onnx_path) != 0:
if model_type == "jina-clip-v1-vision":
load_output = rknn.load_onnx(
model=onnx_path,
inputs=["pixel_values"],
input_size_list=[[1, 3, 224, 224]],
)
else:
load_output = rknn.load_onnx(model=onnx_path)
if load_output != 0:
logger.error("Failed to load ONNX model")
return False
@@ -265,7 +314,7 @@ def is_lock_stale(lock_file_path: Path, max_age: int = 600) -> bool:
def wait_for_conversion_completion(
rknn_path: Path, lock_file_path: Path, timeout: int = 300
model_type: str, rknn_path: Path, lock_file_path: Path, timeout: int = 300
) -> bool:
"""
Wait for another process to complete the conversion.
@@ -307,7 +356,7 @@ def wait_for_conversion_completion(
# Check if RKNN file appeared while waiting
if rknn_path.exists():
logger.info(f"RKNN model appeared while waiting: {rknn_path}")
return str(rknn_path)
return True
# Convert ONNX to RKNN
logger.info(
@@ -320,12 +369,12 @@ def wait_for_conversion_completion(
if onnx_path.exists():
if convert_onnx_to_rknn(
str(onnx_path), str(rknn_path), "yolo-generic", False
str(onnx_path), str(rknn_path), model_type, False
):
return str(rknn_path)
return True
logger.error("Failed to convert model after stale lock cleanup")
return None
return False
finally:
release_conversion_lock(lock_file_path)
@@ -338,7 +387,7 @@ def wait_for_conversion_completion(
def auto_convert_model(
model_path: str, model_type: str, quantization: bool = False
model_path: str, model_type: str | None = None, quantization: bool = False
) -> Optional[str]:
"""
Automatically convert a model to RKNN format if needed.
@@ -377,6 +426,9 @@ def auto_convert_model(
logger.info(f"Converting {model_path} to RKNN format...")
rknn_path.parent.mkdir(parents=True, exist_ok=True)
if not model_type:
model_type = get_rknn_model_type(base_path)
if convert_onnx_to_rknn(
str(base_path), str(rknn_path), model_type, quantization
):
@@ -392,7 +444,10 @@ def auto_convert_model(
f"Another process is converting {model_path}, waiting for completion..."
)
if wait_for_conversion_completion(rknn_path, lock_file_path):
if not model_type:
model_type = get_rknn_model_type(base_path)
if wait_for_conversion_completion(model_type, rknn_path, lock_file_path):
return str(rknn_path)
else:
logger.error(f"Timeout waiting for conversion of {model_path}")