main.py
 1  import cv2
 2  import os
 3  import datetime
 4  from detector import detect
 5  from config import SAVE_IMAGES, VIDEO_SOURCE
 6  
 7  os.makedirs("captures", exist_ok=True)
 8  os.makedirs("logs", exist_ok=True)
 9  
10  cap = cv2.VideoCapture(VIDEO_SOURCE)
11  
12  vehicle_count = 0
13  
14  print("AI Traffic Monitoring System Started... Press Q to Exit")
15  
16  while True:
17      ret, frame = cap.read()
18      if not ret:
19          break
20  
21      detections = detect(frame)
22  
23      for label, conf, x1, y1, x2, y2 in detections:
24          vehicle_count += 1
25  
26          cv2.rectangle(frame, (x1, y1), (x2, y2), (0,255,0), 2)
27          cv2.putText(frame, f"{label} {conf:.2f}", (x1, y1-10),
28                      cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2)
29  
30          if SAVE_IMAGES:
31              filename = f"captures/{label}_{vehicle_count}.jpg"
32              cv2.imwrite(filename, frame)
33  
34          with open("logs/traffic_log.txt", "a") as f:
35              f.write(f"{datetime.datetime.now()} - {label}\n")
36  
37      cv2.putText(frame, f"Vehicle Count: {vehicle_count}", (20, 40),
38                  cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 3)
39  
40      cv2.imshow("AI Traffic Monitoring", frame)
41  
42      if cv2.waitKey(1) & 0xFF == ord('q'):
43          break
44  
45  cap.release()
46  cv2.destroyAllWindows()