File size: 3,778 Bytes
680620c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import streamlit as st
from PIL import Image
import numpy as np
import cv2
from ultralytics import YOLO
# Load your PPE detection model
model = YOLO('ppe-kit.pt') # Replace with your PPE detection model path
# Set page configuration
st.set_page_config(
page_title="PPE Kit Detection",
page_icon="🦺",
layout="wide",
initial_sidebar_state="expanded",
)
# Sidebar contents
st.sidebar.title("🦺 PPE Kit Detection")
st.sidebar.write("Upload an image to detect PPE kits.")
# Upload image
uploaded_file = st.sidebar.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
# Main content
st.title("🦺 PPE Kit Detection")
st.write(
"""
This application uses a state-of-the-art YOLO model to detect PPE kits and safety measures in images.
It identifies items like helmets, masks, safety vests, and more. Upload an image to see the detection in action.
"""
)
if uploaded_file is not None:
# Display uploaded image
image = Image.open(uploaded_file)
# Create two columns for displaying the image
col1, col2 = st.columns(2)
with col1:
st.header("Uploaded Image")
st.image(image, caption='Uploaded Image', use_column_width=True)
with col2:
st.header("Detected PPE and Safety Items")
# Convert the image to OpenCV format
image_cv = np.array(image)
image_cv = image_cv[:, :, ::-1].copy() # Convert RGB to BGR
# Run inference
results = model(image_cv)
# Process the results
boxes = results[0].boxes
labels = {
0: 'Helmet', 1: 'Mask', 2: 'NO-Helmet', 3: 'NO-Mask', 4: 'NO-Safety Vest',
5: 'Person', 6: 'Safety Cone', 7: 'Safety Vest', 8: 'machinery', 9: 'vehicle'
}
detected_items = [] # List to store detected items and accuracy
# Define a set of class IDs to skip (machinery, vehicle, and person)
skip_classes = {5, 8, 9} # Person, machinery, vehicle
# Draw bounding boxes and labels on the image
for box in boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0])
confidence = box.conf[0]
class_id = int(box.cls[0])
label = labels.get(class_id, 'Unknown')
# Skip drawing and storing items if they belong to the skip_classes set
if class_id in skip_classes:
continue
# Store detected items with accuracy
detected_items.append(f"{label} ({confidence * 100:.2f}%)")
# Draw the box
cv2.rectangle(image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)
# Put the label
cv2.putText(image_cv, f"{label} {confidence:.2f}", (x1, y1 - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
# Convert the image back to RGB
image_cv = image_cv[:, :, ::-1].copy()
st.image(image_cv, caption='Detected PPE and Safety Items', use_column_width=True)
# Centered detected items below the image
st.markdown("<h2 style='text-align: center;'>Detected Items with Accuracy:</h2>", unsafe_allow_html=True)
# Display the detected items in a centered fashion
for item in detected_items:
st.markdown(f"<h4 style='text-align: center;'>{item}</h4>", unsafe_allow_html=True)
st.success("Detection complete!")
else:
st.info("Please upload an image to start detection.")
# Footer
st.sidebar.markdown("### About")
st.sidebar.info(
"This app uses a YOLO model to detect PPE kits like helmets, masks, safety vests, and other safety-related items. "
"Upload an image and the model will identify these items in the image."
)
|