add sam2 yolo auto annotation

2026-02-04 15:29:36 +07:00
parent 7e56948ece
commit 5a951d8812
2061 changed files with 316473 additions and 0 deletions
@@ -0,0 +1,118 @@
+"""
+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