Spaces:
Sleeping
Sleeping
File size: 3,943 Bytes
2353f19 385665b 2353f19 385665b 2353f19 |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
import streamlit as st
from ultralytics import YOLO
from PIL import Image
import torchvision.transforms as transforms
import base64
import cv2
import numpy as np
# Set Streamlit Page Configuration
st.set_page_config(
page_title="PPE Detect",
page_icon="logo/logo.png",
layout="centered"
)
# Cache the YOLO model to optimize performance
@st.cache_resource()
def load_model():
return YOLO("model/best.pt") # Ensure correct model path
model = load_model()
# Define image transformation pipeline
transform = transforms.Compose([
transforms.Resize((640, 640)),
transforms.ToTensor()
])
# Function to perform PPE detection on images
def predict_ppe(image: Image.Image):
try:
image_tensor = transform(image).unsqueeze(0) # Add batch dimension
results = model.predict(image_tensor)
output_image = results[0].plot() # Overlay predictions
return Image.fromarray(output_image)
except Exception as e:
st.error(f"Prediction Error: {e}")
return None
# Function to encode image to base64 for embedding
def get_base64_image(image_path):
try:
with open(image_path, "rb") as img_file:
return base64.b64encode(img_file.read()).decode()
except FileNotFoundError:
return None
# Function for real-time PPE detection using webcam
def live_ppe_detection():
st.sidebar.write("Starting live detection...")
cap = cv2.VideoCapture(0)
if not cap.isOpened():
st.sidebar.error("Error: Could not open webcam.")
return
stframe = st.empty()
stop_button = st.sidebar.button("Stop Live Detection", key="stop_button")
while cap.isOpened():
ret, frame = cap.read()
if not ret:
st.sidebar.error("Failed to capture video frame.")
break
results = model.predict(frame)
output_frame = results[0].plot()
stframe.image(output_frame, channels="BGR")
if stop_button:
break
cap.release()
cv2.destroyAllWindows()
# Display logo
image_base64 = get_base64_image("logo/logo.png")
if image_base64:
st.markdown(
f'<div style="text-align: center;"><img src="data:image/png;base64,{image_base64}" width="100"></div>',
unsafe_allow_html=True
)
# UI Customization
st.markdown("""
<style>
[data-testid="stSidebar"] { background-color: #1E1E2F; }
[data-testid="stSidebar"] h1, [data-testid="stSidebar"] h2 { color: white; }
h1 { text-align: center; font-size: 36px; font-weight: bold; color: #2C3E50; }
div.stButton > button { background-color: #3498DB; color: white; font-weight: bold; }
div.stButton > button:hover { background-color: #2980B9; }
</style>
""", unsafe_allow_html=True)
# Sidebar - File Upload
st.sidebar.header("📤 Upload an Image")
uploaded_file = st.sidebar.file_uploader("Drag and drop or browse", type=['jpg', 'png', 'jpeg'])
# Sidebar - Live Predictions
st.sidebar.header("📡 Live Predictions")
if st.sidebar.button("Start Live Detection", key="start_button"):
live_ppe_detection()
# Main Page
st.title("PPE Detect")
st.markdown("<p style='text-align: center;'>Detect personal protective equipment (PPE) in images.</p>", unsafe_allow_html=True)
if uploaded_file:
image = Image.open(uploaded_file).convert("RGB")
col1, col2 = st.columns(2)
with col1:
st.image(image, caption="📷 Uploaded Image", use_container_width=True)
if st.sidebar.button("🔍 Predict PPE", key="predict_button"):
detected_image = predict_ppe(image)
if detected_image:
with col2:
st.image(detected_image, caption="🎯 PPE Detection Result", use_container_width=True)
else:
st.error("Detection failed. Please try again.")
st.info("This app uses **YOLO** for PPE detection. Upload an image or start live detection to get started.")
|