Spaces:
Sleeping
Sleeping
Upload 3 files
Browse files- hdv.py +36 -0
- helmet_detector.py +43 -0
- 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
|