Helmet / app_streamlit.py
NightPrince's picture
Update app_streamlit.py
6811ebf verified
import streamlit as st
from PIL import Image
from ultralytics import YOLO
import numpy as np
from huggingface_hub import hf_hub_download
# Download the model from Hugging Face Hub
@st.cache_resource
def load_model():
model_path = hf_hub_download(repo_id="NightPrince/Yolo-Helmet", filename="best.pt")
model = YOLO(model_path)
return model
model = load_model()
st.title("🚨 Helmet Detection with YOLOv8 🚧")
st.write("Upload an image, and the model will detect if the driver is wearing a helmet.")
uploaded_file = st.file_uploader("📸 Upload Image", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
image = Image.open(uploaded_file)
st.image(image, caption='📥 Uploaded Image', use_column_width=True)
st.write("🔍 Running YOLOv8 Detection...")
results = model.predict(image)
result_img = results[0].plot()
st.image(result_img, caption='✅ Detection Result', use_column_width=True)
st.write("📋 Detection Details:")
st.dataframe(results[0].to_pandas())