Spaces:
Running on Zero
Running on Zero
| """ | |
| Phase 5 - Emergency vehicle detection. | |
| Uses emergency.pt to flag emergency vs non-emergency vehicles. | |
| """ | |
| from ultralytics import YOLO | |
| import sys | |
| import os | |
| EMERGENCY_MODEL_PATH = r"E:\Traffic_AI\models\emergency.pt" | |
| OUTPUT_DIR = r"E:\Traffic_AI\outputs" | |
| os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| EMERGENCY_CLASS_ID = 0 # 'emergency_vehicle' | |
| def detect_emergency(source: str, mode: str): | |
| print(f"\n--- Phase 5 Emergency Vehicle Detection: {source} ---") | |
| model = YOLO(EMERGENCY_MODEL_PATH) | |
| is_video = mode == "video" | |
| results = model.predict( | |
| source=source, | |
| conf=0.30, | |
| device=0, | |
| save=True, | |
| project=OUTPUT_DIR, | |
| name=f"{mode}_emergency", | |
| stream=is_video, | |
| ) | |
| frame_count = 0 | |
| emergency_count = 0 | |
| for r in results: | |
| frame_count += 1 | |
| frame_emergency = [] | |
| for box in r.boxes: | |
| cls_id = int(box.cls) | |
| label = model.names[cls_id] | |
| conf = round(float(box.conf), 2) | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 | |
| if cls_id == EMERGENCY_CLASS_ID: | |
| emergency_count += 1 | |
| frame_emergency.append((label, conf, cx, cy, x1, y1, x2, y2)) | |
| if not is_video and frame_emergency: | |
| print(f"EMERGENCY VEHICLES DETECTED: {len(frame_emergency)}") | |
| for label, conf, cx, cy, x1, y1, x2, y2 in frame_emergency: | |
| print(f" {label:20s} conf={conf:.2f} center=({cx},{cy}) box=[{x1},{y1},{x2},{y2}]") | |
| elif not is_video: | |
| print("No emergency vehicles detected.") | |
| if is_video and frame_count % 30 == 0: | |
| if frame_emergency: | |
| print(f" Frame {frame_count}: {len(frame_emergency)} EMERGENCY vehicle(s) detected!") | |
| for label, conf, cx, cy, x1, y1, x2, y2 in frame_emergency: | |
| print(f" conf={conf:.2f} center=({cx},{cy})") | |
| else: | |
| print(f" Frame {frame_count}: no emergency vehicles") | |
| print(f"\nTotal frames processed: {frame_count}") | |
| print(f"Total emergency detections: {emergency_count}") | |
| print(f"Output saved to: {OUTPUT_DIR}\\{mode}_emergency") | |
| if __name__ == "__main__": | |
| if len(sys.argv) < 3: | |
| print("Usage:") | |
| print(" Image: python detect_emergency.py image <path>") | |
| print(" Video: python detect_emergency.py video <path>") | |
| sys.exit(1) | |
| mode = sys.argv[1].lower() | |
| path = sys.argv[2] | |
| if not os.path.exists(path): | |
| print(f"ERROR: File not found: {path}") | |
| sys.exit(1) | |
| if mode not in ("image", "video"): | |
| print("ERROR: mode must be 'image' or 'video'") | |
| sys.exit(1) | |
| detect_emergency(path, mode) | |