root007x/Worker_Helmet_Detection

1import gradio as gr2import cv23import os4import numpy as np5from ultralytics import YOLO6 7 8model = YOLO("final_model.pt")9 10classes = ['head', 'helmet', 'person']11 12 13def show_pred_image(image_path):14    15    image = cv2.resize(image_path, (640,640))16    # Predict using the YOLO model17    outputs = model.predict(source=image)18    results = outputs[0]  # Get the first result19    20    boxes = results.boxes.xyxy.cpu().numpy()  # Bounding boxes21    confidences = results.boxes.conf.cpu().numpy()  # Confidence scores22    labels = results.boxes.cls.cpu().numpy()  # Class labels23    24    for i, box in enumerate(boxes):25        x1, y1, x2, y2 = map(int, box)  # Coordinates of the bounding box26        27        # Draw the bounding box on the image28        cv2.rectangle(image, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=2, lineType=cv2.LINE_AA)29        30        # Prepare the label text (class + confidence)31        label = f"{classes[int(labels[i])]}: {confidences[i]:.2f}"32        33        # Put the label text above the bounding box34        cv2.putText(35            image,36            label,37            (x1, y1 - 10),  # Place the text slightly above the top-left corner of the box38            cv2.FONT_HERSHEY_SIMPLEX,39            0.6,  # Font scale40            color = (0, 255, 0),41            thickness = 1,  42            lineType = cv2.LINE_AA43        )44    45    return image46 47 48example_list = [["examples/" + example] for example in os.listdir("examples")]49 50inputs_image = gr.Image(label = "Input Image")51outputs_image = gr.Image(label = "Output Image")52 53demo = gr.Interface(54    fn = show_pred_image,55    inputs = inputs_image,56    outputs = outputs_image,57    title = "Worker Helmet Detection",58    examples = example_list59)60 61demo.launch()