119 lines
4.4 KiB
Python
119 lines
4.4 KiB
Python
"""
|
|
Object detection backends.
|
|
"""
|
|
|
|
from typing import Dict, Any, Optional
|
|
import logging
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def create_detector(config: Dict[str, Any]):
|
|
"""
|
|
Create detector based on configuration.
|
|
|
|
Args:
|
|
config: Detector configuration dict
|
|
|
|
Returns:
|
|
Detector instance
|
|
"""
|
|
detector_type = config.get('type', 'yolo').lower()
|
|
model_path = config.get('model_path', 'models/yolov9t.pt')
|
|
|
|
logger.info(f"Creating detector: type={detector_type}, model={model_path}")
|
|
|
|
# Try RKNN first if specified
|
|
if detector_type == 'rknn':
|
|
try:
|
|
from .rknn_detector import RKNNDetector
|
|
|
|
rknn_config = config.get('rknn', {})
|
|
detector = RKNNDetector(
|
|
model_path=model_path,
|
|
target_platform=rknn_config.get('target_platform', 'rk3588'),
|
|
core_mask=rknn_config.get('core_mask', 7),
|
|
input_size=tuple(config.get('input_size', [640, 640])),
|
|
conf_threshold=config.get('conf_threshold', 0.25),
|
|
nms_threshold=config.get('nms_threshold', 0.45),
|
|
)
|
|
|
|
if detector.load_model():
|
|
logger.info("RKNN detector initialized successfully")
|
|
return detector
|
|
else:
|
|
logger.warning("RKNN detector failed to load, trying fallback")
|
|
|
|
except ImportError as e:
|
|
logger.warning(f"RKNN not available: {e}")
|
|
except Exception as e:
|
|
logger.warning(f"RKNN initialization failed: {e}")
|
|
|
|
# Try ONNX if specified or as fallback
|
|
if detector_type == 'onnx' or (detector_type == 'rknn' and config.get('fallback', {}).get('enabled', True)):
|
|
fallback_config = config.get('fallback', {})
|
|
onnx_config = config.get('onnx', {})
|
|
|
|
if fallback_config.get('type') == 'onnx' or detector_type == 'onnx':
|
|
try:
|
|
from .onnx_detector import ONNXDetector
|
|
|
|
# Determine model path
|
|
onnx_model_path = model_path
|
|
if model_path.endswith('.rknn'):
|
|
onnx_model_path = model_path.replace('.rknn', '.onnx')
|
|
elif model_path.endswith('.pt'):
|
|
onnx_model_path = model_path.replace('.pt', '.onnx')
|
|
|
|
# Get device from onnx config or fallback config
|
|
device = onnx_config.get('device') or fallback_config.get('device', 'cpu')
|
|
|
|
detector = ONNXDetector(
|
|
model_path=onnx_model_path,
|
|
input_size=tuple(config.get('input_size', [640, 640])),
|
|
conf_threshold=config.get('conf_threshold', 0.25),
|
|
nms_threshold=config.get('nms_threshold', 0.45),
|
|
device=device,
|
|
num_threads=onnx_config.get('num_threads', 0),
|
|
optimization_level=onnx_config.get('optimization_level', 'all'),
|
|
)
|
|
|
|
if detector.load_model():
|
|
logger.info("ONNX detector initialized successfully")
|
|
return detector
|
|
|
|
except ImportError as e:
|
|
logger.warning(f"ONNX runtime not available: {e}")
|
|
logger.info("Install with: pip install onnxruntime")
|
|
except Exception as e:
|
|
logger.warning(f"ONNX initialization failed: {e}")
|
|
|
|
# Use Ultralytics YOLO as default/fallback
|
|
try:
|
|
from .yolo_detector import YOLODetector
|
|
|
|
fallback_config = config.get('fallback', {})
|
|
device = fallback_config.get('device', 'cpu')
|
|
|
|
# Adjust model path
|
|
if model_path.endswith('.rknn'):
|
|
model_path = model_path.replace('.rknn', '.pt')
|
|
elif model_path.endswith('.onnx'):
|
|
model_path = model_path.replace('.onnx', '.pt')
|
|
|
|
detector = YOLODetector(
|
|
model_path=model_path,
|
|
conf_threshold=config.get('conf_threshold', 0.25),
|
|
nms_threshold=config.get('nms_threshold', 0.45),
|
|
device=device,
|
|
)
|
|
|
|
if detector.load_model():
|
|
logger.info(f"YOLO detector initialized on {device}")
|
|
return detector
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to initialize any detector: {e}")
|
|
|
|
return None
|