import streamlit as st
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image

# Define repository and file path
repo_id = "krishnamishra8848/Face_Mask_Detection"
filename = "best.pt"  # File name in your Hugging Face repo

# Download the model file
model_path = hf_hub_download(repo_id=repo_id, filename=filename)

# Load the YOLO model
model = YOLO(model_path)

# Streamlit UI
st.title("Face Mask Detection with YOLOv8")
st.write("Upload an image to detect face masks.")

# File upload
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file:
    # Load image
    image = Image.open(uploaded_file)
    image_np = np.array(image)

    # Display "Running inference..." in red
    placeholder = st.empty()
    placeholder.markdown('<h3 style="color: red;">Running inference...</h3>', unsafe_allow_html=True)

    # Run inference
    results = model.predict(source=image_np, conf=0.5)

    # Annotate image
    annotated_image = None
    for result in results:
        annotated_image = result.plot()

    # Convert annotated image for Streamlit
    if annotated_image is not None:
        annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
        placeholder.empty()  # Remove the "Running inference..." message
        st.image(annotated_image_rgb, caption="Prediction Results", use_container_width=True)