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())