NarendraChari commited on
Commit
a8917eb
·
verified ·
1 Parent(s): b8130ce

Upload 3 files

Browse files
Files changed (3) hide show
  1. hdv.py +36 -0
  2. helmet_detector.py +43 -0
  3. requiremets.txt +2 -0
hdv.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from ultralytics import YOLO
3
+
4
+ # Initialize video capture
5
+ video_path = "Media/traffic.mp4"
6
+ cap = cv2.VideoCapture(video_path)
7
+
8
+ # Load YOLO model with custom weights
9
+ model = YOLO("Weights/best.pt")
10
+
11
+ # Define class names
12
+ classNames = ['With Helmet', 'Without Helmet']
13
+
14
+ # For the use of Webcam
15
+ # Open the webcam (use 0 for the default camera, or 1, 2, etc. for additional cameras)
16
+ # cap = cv2.VideoCapture(0)
17
+
18
+ while True:
19
+ success, img = cap.read()
20
+ results = model(img, stream=True)
21
+ for r in results:
22
+ boxes = r.boxes
23
+ for box in boxes:
24
+ x1, y1, x2, y2 = box.xyxy[0]
25
+ x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
26
+
27
+ w, h = x2 - x1, y2 - y1
28
+ cvzone.cornerRect(img, (x1, y1, w, h))
29
+ conf = math.ceil((box.conf[0] * 100)) / 100
30
+ cls = int(box.cls[0])
31
+
32
+ cvzone.putTextRect(img, f'{classNames[cls]} {conf}', (max(0, x1), max(35, y1)), scale=1, thickness=1)
33
+
34
+ cv2.imshow("Image", img)
35
+ if cv2.waitKey(1) & 0xFF == ord('q'):
36
+ break
helmet_detector.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import math
3
+ import cvzone
4
+ from ultralytics import YOLO
5
+
6
+ # Load YOLO model with custom weights
7
+ yolo_model = YOLO("Weights/best.pt")
8
+
9
+ # Define class names
10
+ class_labels = ['With Helmet', 'Without Helmet']
11
+
12
+ # Load the image
13
+ image_path = "Media/riders_1.jpg"
14
+ img = cv2.imread(image_path)
15
+
16
+ # Perform object detection
17
+ results = yolo_model(img)
18
+
19
+ # Loop through the detections and draw bounding boxes
20
+ for r in results:
21
+ boxes = r.boxes
22
+ for box in boxes:
23
+ x1, y1, x2, y2 = box.xyxy[0]
24
+ x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
25
+
26
+ w, h = x2 - x1, y2 - y1
27
+ cvzone.cornerRect(img, (x1, y1, w, h))
28
+ conf = math.ceil((box.conf[0] * 100)) / 100
29
+ cls = int(box.cls[0])
30
+
31
+ if conf > 0.1:
32
+ cvzone.putTextRect(img, f'{class_labels[cls]} {conf}', (x1, y1 - 10), scale=0.8, thickness=1, colorR=(255, 0, 0))
33
+
34
+ # Display the image with detections
35
+ cv2.imshow("Image", img)
36
+
37
+ # Close window when 'q' button is pressed
38
+ while True:
39
+ if cv2.waitKey(1) & 0xFF == ord('q'):
40
+ break
41
+
42
+ cv2.destroyAllWindows()
43
+ cv2.waitKey(1)
requiremets.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ yolov9pip==0.0.4
2
+ huggingface_hub