ObjectDetection.py

import cv2 as cv
import cvzone
import math
from ultralytics import YOLO

# Initialize video capture and model
cap = cv.VideoCapture("8488067.mp4")  
model = YOLO("ppe.pt")

# Define class names for detection
classNames = ['Hardhat', 'Mask', 'NO-Hardhat', 'NO-Mask', 'NO-Safety Vest', 'Person', 'Safety Cone',
              'Safety Vest', 'machinery', 'vehicle']

# Initializing bounding box color
myColor = (0, 0, 255)



while True:
    # Read frame from video
    success, img = cap.read()
    results = model(img, stream=True)


    for r in results:
        boxes = r.boxes

        for box in boxes:
             # Extract bounding box coordinates and confidence score
            x1, y1, x2, y2 = box.xyxy[0]
            x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
            w, h = x2 - x1, y2 - y1
            # Confidence
            conf = math.ceil((box.conf[0] * 100)) / 100
            # Class Name
            cls = int(box.cls[0])
            currentClass = classNames[cls]
            

            # Filter for relevant classes based on confidence threshold
            if conf>0.5:
                if currentClass in ["NO-Hardhat", "NO-Safety Vest", "NO-Mask"]:
                    myColor = (0, 0,255)
                elif currentClass in ["Hardhat", "Safety Vest", "Mask"]:
                    myColor =(0,255,0)
                else:
                    myColor = (255, 0, 0)

                # Display Class Name, and Draw bounding box on the image
                cvzone.putTextRect(img, f'{classNames[cls]} {conf}',
                                   (max(0, x1), max(35, y1)), scale=1, thickness=1,colorB=myColor,
                                   colorT=(255,255,255),colorR=myColor, offset=5)
                cv.rectangle(img, (x1, y1), (x2, y2), myColor, 3)
 
    cv.imshow("Image", img)
    # Break loop on key press
    if cv.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv.destroyAllWindows()