import os import sys import time import cv2 import numpy as np from pathlib import Path from typing import Dict, List, Tuple, Optional # Add parent directory to path for imports current_dir = Path(__file__).parent.parent.parent sys.path.append(str(current_dir)) # Import OpenVINO modules from detection_openvino import OpenVINOVehicleDetector from red_light_violation_pipeline import RedLightViolationPipeline # Import from our utils package from utils.helpers import bbox_iou class ModelManager: """ Manages OpenVINO models for traffic detection and violation monitoring. Only uses RedLightViolationPipeline for all violation/crosswalk/traffic light logic. """ def __init__(self, config_file: str = None, tracker=None): """ Initialize model manager with configuration. Args: config_file: Path to JSON configuration file tracker: (Optional) External tracker instance (e.g., DeepSortVehicleTracker singleton) """ self.config = self._load_config(config_file) self.detector = None self.violation_pipeline = None # Use RedLightViolationPipeline only self.tracker = tracker self._initialize_models() def _load_config(self, config_file: Optional[str]) -> Dict: """ Load configuration from file or use defaults. Args: config_file: Path to JSON configuration file Returns: Configuration dictionary """ import json default_config = { "detection": { "confidence_threshold": 0.3, "enable_ocr": True, "enable_tracking": True, "model_path": None }, "violations": { "red_light_grace_period": 2.0, "stop_sign_duration": 2.0, "speed_tolerance": 5 }, "display": { "max_display_width": 800, "show_confidence": True, "show_labels": True, "show_license_plates": True }, "performance": { "max_history_frames": 1000, "cleanup_interval": 3600 } } if config_file and os.path.exists(config_file): try: with open(config_file, 'r') as f: loaded_config = json.load(f) # Merge with defaults (preserving loaded values) for section in default_config: if section in loaded_config: default_config[section].update(loaded_config[section]) except Exception as e: print(f"Error loading config: {e}") return default_config def _initialize_models(self): """Initialize OpenVINO detection and violation models.""" try: # Find best model path model_path = self.config["detection"].get("model_path") if not model_path or not os.path.exists(model_path): model_path = self._find_best_model_path() if not model_path: print("❌ No model found") return # Initialize detector print(f"✅ Initializing OpenVINO detector with model: {model_path}") device = self.config["detection"].get("device", "AUTO") print(f"✅ Using inference device: {device}") self.detector = OpenVINOVehicleDetector( model_path=model_path, device=device, confidence_threshold=self.config["detection"]["confidence_threshold"] ) # Use only RedLightViolationPipeline for violation/crosswalk/traffic light logic self.violation_pipeline = RedLightViolationPipeline(debug=True) print("✅ Red light violation pipeline initialized (all other violation logic removed)") # Only initialize tracker if not provided if self.tracker is None and self.config["detection"]["enable_tracking"]: try: from controllers.bytetrack_tracker import ByteTrackVehicleTracker self.tracker = ByteTrackVehicleTracker() print("✅ ByteTrack tracker initialized (internal)") except ImportError: print("⚠️ ByteTrack not available") self.tracker = None elif self.tracker is not None: print("✅ Using external DeepSORT tracker instance") print("✅ Models initialized successfully") except Exception as e: print(f"❌ Error initializing models: {e}") import traceback traceback.print_exc() def _find_best_model_path(self, base_model_name: str = None) -> Optional[str]: if base_model_name is None: device = self.config["detection"].get("device", "AUTO") if device == "CPU" or device == "AUTO": # Use yolo11n for CPU - faster, lighter model base_model_name = "yolo11n" print(f"🔍 Device is {device}, selecting {base_model_name} model (optimized for CPU)") else: # Use yolo11x for GPU - larger model with better accuracy base_model_name = "yolo11x" print(f"🔍 Device is {device}, selecting {base_model_name} model (optimized for GPU)") # Check if the openvino_models directory exists in the current working directory cwd_openvino_dir = Path.cwd() / "openvino_models" if cwd_openvino_dir.exists(): direct_path = cwd_openvino_dir / f"{base_model_name}.xml" if direct_path.exists(): print(f"✅ Found model directly in CWD: {direct_path}") return str(direct_path.absolute()) # Check for absolute path to openvino_models (this is the most reliable) absolute_openvino_dir = Path("D:/Downloads/finale6/khatam/openvino_models") if absolute_openvino_dir.exists(): direct_path = absolute_openvino_dir / f"{base_model_name}.xml" if direct_path.exists(): print(f"✅ Found model at absolute path: {direct_path}") return str(direct_path.absolute()) # Try relative to the model_manager.py file openvino_models_dir = Path(__file__).parent.parent.parent / "openvino_models" direct_path = openvino_models_dir / f"{base_model_name}.xml" if direct_path.exists(): print(f"✅ Found model in app directory: {direct_path}") return str(direct_path.absolute()) # Check for model in folder structure within openvino_models subfolder_path = openvino_models_dir / f"{base_model_name}_openvino_model" / f"{base_model_name}.xml" if subfolder_path.exists(): print(f"✅ Found model in subfolder: {subfolder_path}") return str(subfolder_path.absolute()) # Try other common locations search_dirs = [ ".", "..", "../models", "../rcb", "../openvino_models", f"../{base_model_name}_openvino_model", "../..", # Go up to project root "../../openvino_models", # Project root / openvino_models ] model_extensions = [ (f"{base_model_name}.xml", "OpenVINO IR direct"), (f"{base_model_name}_openvino_model/{base_model_name}.xml", "OpenVINO IR"), (f"{base_model_name}.pt", "PyTorch"), ] for search_dir in search_dirs: search_path = Path(__file__).parent.parent / search_dir if not search_path.exists(): continue for model_file, model_type in model_extensions: model_path = search_path / model_file if model_path.exists(): print(f"✅ Found {model_type} model: {model_path}") return str(model_path.absolute()) print(f"❌ No model found for {base_model_name}") return None def detect(self, frame: np.ndarray) -> List[Dict]: """ Detect objects in frame. Args: frame: Input video frame Returns: List of detection dictionaries """ if self.detector is None: print("WARNING: No detector available") return [] try: # Use a lower confidence threshold for better visibility base_conf_threshold = self.config["detection"].get("confidence_threshold", 0.5) conf_threshold = max(0.15, base_conf_threshold) # Lowered to 0.15 for traffic lights detections = self.detector.detect_vehicles(frame, conf_threshold=conf_threshold) # Try to find traffic lights with even lower confidence if none found traffic_light_found = any(det.get('class_name') == 'traffic light' for det in detections) if not traffic_light_found: print("⚠️ No traffic lights detected with normal confidence, trying lower threshold...") try: low_conf_detections = self.detector.detect_vehicles(frame, conf_threshold=0.05) for det in low_conf_detections: if det.get('class_name') == 'traffic light' and det not in detections: print(f"🚦 Adding low confidence traffic light: conf={det['confidence']:.3f}") detections.append(det) except Exception as e: print(f"❌ Error trying low confidence detection: {e}") # Enhance traffic light detection using the same utilities as qt_app_pyside from utils.traffic_light_utils import detect_traffic_light_color, ensure_traffic_light_color for det in detections: if det.get('class_id') == 9 or det.get('class_name') == 'traffic light': try: bbox = det['bbox'] light_info = detect_traffic_light_color(frame, bbox) if light_info.get("color", "unknown") == "unknown": light_info = ensure_traffic_light_color(frame, bbox) det['traffic_light_color'] = light_info print(f"🚦 Enhanced Traffic Light Detection: {light_info}") except Exception as e: print(f"❌ Error in enhanced traffic light detection: {e}") # Ensure all detections have valid class_name and confidence for det in detections: if det.get('class_name') is None: det['class_name'] = 'object' if det.get('confidence') is None: det['confidence'] = 0.0 # Add debug output if detections: print(f"DEBUG: Detected {len(detections)} objects: " + ", ".join([f"{d['class_name']} ({d['confidence']:.2f})" for d in detections[:3]])) # Print bounding box coordinates of first detection if len(detections) > 0: print(f"DEBUG: First detection bbox: {detections[0]['bbox']}") else: print("DEBUG: No detections in this frame") return detections except Exception as e: print(f"❌ Detection error: {e}") import traceback traceback.print_exc() return [] def update_tracking(self, detections: List[Dict], frame: np.ndarray) -> List[Dict]: """ Update tracking information for detections. Args: detections: List of detections frame: Current video frame Returns: Updated list of detections with tracking info """ if not self.tracker or not detections: # Fallback: assign temporary IDs if no tracker for idx, det in enumerate(detections): det['id'] = idx if det.get('class_name') is None: det['class_name'] = 'object' if det.get('confidence') is None: det['confidence'] = 0.0 return detections try: tracker_dets = [] det_map = [] # Keep mapping to original detection for det in detections: bbox = det['bbox'] if len(bbox) < 4: continue x1, y1, x2, y2 = bbox w = x2 - x1 h = y2 - y1 if w <= 0 or h <= 0: continue conf = det.get('confidence', 0.0) class_name = det.get('class_name', 'object') tracker_dets.append(([x1, y1, w, h], conf, class_name)) det_map.append(det) # Update tracks output = [] if tracker_dets: tracks = self.tracker.update_tracks(tracker_dets, frame=frame) for i, track in enumerate(tracks): # FIXED: Handle both object-style tracks (with methods) and dict-style tracks # First check if track is confirmed (handle both dict and object styles) is_confirmed = True # Default to True for dict-style tracks if hasattr(track, 'is_confirmed') and callable(getattr(track, 'is_confirmed')): is_confirmed = track.is_confirmed() if not is_confirmed: continue # Get track_id (handle both dict and object styles) if hasattr(track, 'track_id'): track_id = track.track_id elif isinstance(track, dict) and 'id' in track: track_id = track['id'] else: print(f"Warning: Track has no ID, skipping: {track}") continue # Get bounding box (handle both dict and object styles) if hasattr(track, 'to_ltrb') and callable(getattr(track, 'to_ltrb')): ltrb = track.to_ltrb() elif isinstance(track, dict) and 'bbox' in track: ltrb = track['bbox'] # Assume bbox is already in [x1,y1,x2,y2] format else: print(f"Warning: Track has no bbox, skipping: {track}") continue # Try to match track to detection by index (DeepSORT returns tracks in same order as input detections) # If not, fallback to previous logic matched_class = 'object' matched_conf = 0.0 if hasattr(track, 'det_index') and track.det_index is not None and track.det_index < len(det_map): matched_class = det_map[track.det_index].get('class_name', 'object') matched_conf = det_map[track.det_index].get('confidence', 0.0) else: # Try to match by IoU if possible best_iou = 0 for det in det_map: db = det['bbox'] iou = self._bbox_iou([int(ltrb[0]), int(ltrb[1]), int(ltrb[2]), int(ltrb[3])], db) if iou > best_iou: best_iou = iou matched_class = det.get('class_name', 'object') matched_conf = det.get('confidence', 0.0) if matched_class is None: matched_class = 'object' if matched_conf is None: matched_conf = 0.0 output.append({ 'bbox': [int(ltrb[0]), int(ltrb[1]), int(ltrb[2]), int(ltrb[3])], 'class_name': matched_class, 'confidence': matched_conf, 'id': track_id }) # Fallback: assign temp IDs if no tracks if not output: for idx, det in enumerate(detections): det['id'] = idx if det.get('class_name') is None: det['class_name'] = 'object' if det.get('confidence') is None: det['confidence'] = 0.0 return detections return output except Exception as e: print(f"❌ Tracking error: {e}") # Fallback: assign temp IDs for idx, det in enumerate(detections): det['id'] = idx if det.get('class_name') is None: det['class_name'] = 'object' if det.get('confidence') is None: det['confidence'] = 0.0 return detections def update_config(self, new_config: Dict): """ Update configuration parameters. Args: new_config: New configuration dictionary """ if not new_config: return # Store old device setting to check if it changed old_device = self.config["detection"].get("device", "AUTO") if "detection" in self.config else "AUTO" # Update configuration for section in new_config: if section in self.config: self.config[section].update(new_config[section]) else: self.config[section] = new_config[section] # Check if device changed - if so, we need to reinitialize models new_device = self.config["detection"].get("device", "AUTO") device_changed = old_device != new_device if device_changed: print(f"📢 Device changed from {old_device} to {new_device}, reinitializing models...") # Reinitialize models with new device self._initialize_models() return # Just update detector confidence threshold if device didn't change if self.detector: conf_thres = self.config["detection"].get("confidence_threshold", 0.5) self.detector.conf_thres = conf_thres def _bbox_iou(self, boxA, boxB): # Compute the intersection over union of two boxes xA = max(boxA[0], boxB[0]) yA = max(boxA[1], boxB[1]) xB = min(boxA[2], boxB[2]) yB = min(boxA[3], boxB[3]) interArea = max(0, xB - xA) * max(0, yB - yA) boxAArea = max(0, boxA[2] - boxA[0]) * max(0, boxA[3] - boxA[1]) boxBArea = max(0, boxB[2] - boxB[0]) * max(0, boxB[3] - boxB[1]) if boxAArea + boxBArea - interArea == 0: return 0.0 iou = interArea / float(boxAArea + boxBArea - interArea) return iou def switch_model(self, target_device: str = None) -> bool: """ Manually switch to a different model based on target device. Args: target_device: Target device ("CPU", "GPU", "AUTO", etc.) Returns: True if switch was successful, False otherwise """ if target_device: old_device = self.config["detection"].get("device", "AUTO") self.config["detection"]["device"] = target_device print(f"🔄 Manual model switch requested: {old_device} → {target_device}") # If detector has a switch_model method, use it if hasattr(self.detector, 'switch_model'): try: success = self.detector.switch_model(device=target_device) if success: print(f"✅ Successfully switched to {target_device} optimized model") # If tracker needs update, reinitialize if device changed if old_device != target_device: self._initialize_models() # Optionally update tracker return True else: print(f"❌ Failed to switch detector to {target_device}") self.config["detection"]["device"] = old_device return False except Exception as e: print(f"❌ Failed to switch model: {e}") self.config["detection"]["device"] = old_device return False else: # Fallback: reinitialize models try: self._initialize_models() print(f"✅ Successfully switched to {target_device} optimized model (fallback)") return True except Exception as e: print(f"❌ Failed to switch model: {e}") self.config["detection"]["device"] = old_device return False return False