Spaces:
Running on Zero
Running on Zero
| """ | |
| Phase 2 - Detect traffic-relevant objects only. | |
| Classes: car, bus, truck, motorcycle, person | |
| """ | |
| from ultralytics import YOLO | |
| import sys | |
| import os | |
| MODEL_PATH = r"E:\Traffic_AI\models\yolo11n.pt" | |
| OUTPUT_DIR = r"E:\Traffic_AI\outputs" | |
| os.makedirs(OUTPUT_DIR, exist_ok=True) | |
| # COCO class IDs for traffic-relevant objects | |
| TRAFFIC_CLASSES = { | |
| 0: "person", | |
| 2: "car", | |
| 3: "motorcycle", | |
| 5: "bus", | |
| 7: "truck", | |
| } | |
| def detect_traffic(source: str, mode: str): | |
| print(f"\n--- Phase 2 Traffic Detection: {source} ---") | |
| model = YOLO(MODEL_PATH) | |
| is_video = mode == "video" | |
| results = model.predict( | |
| source=source, | |
| conf=0.25, | |
| device=0, | |
| classes=list(TRAFFIC_CLASSES.keys()), # filter classes here | |
| save=True, | |
| project=OUTPUT_DIR, | |
| name=f"{mode}_traffic", | |
| stream=is_video, | |
| ) | |
| frame_count = 0 | |
| for r in results: | |
| frame_count += 1 | |
| detections = [] | |
| for box in r.boxes: | |
| cls_id = int(box.cls) | |
| label = TRAFFIC_CLASSES[cls_id] | |
| conf = float(box.conf) | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| detections.append((label, conf, x1, y1, x2, y2)) | |
| if not is_video: | |
| print(f"Detected {len(detections)} traffic objects:") | |
| for label, conf, x1, y1, x2, y2 in detections: | |
| print(f" {label:12s} conf={conf:.2f} box=[{x1},{y1},{x2},{y2}]") | |
| elif frame_count % 30 == 0: | |
| print(f" Frame {frame_count}: {len(detections)} detections") | |
| if is_video: | |
| print(f"Total frames processed: {frame_count}") | |
| print(f"Output saved to: {OUTPUT_DIR}\\{mode}_traffic") | |
| if __name__ == "__main__": | |
| if len(sys.argv) < 3: | |
| print("Usage:") | |
| print(" Image: python detect_traffic.py image <path>") | |
| print(" Video: python detect_traffic.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_traffic(path, mode) | |