475 lines
21 KiB
Python
475 lines
21 KiB
Python
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
|