import streamlit as st
import cv2
import numpy as np
from ultralytics import YOLO
from PIL import Image
import tempfile
import os

# ---------------- PAGE CONFIG ----------------
st.set_page_config(
    page_title="Helmet Detection",
    page_icon="🪖",
    layout="wide"
)

st.title("🪖 Helmet Detection System")
st.write("Upload an image or video to detect helmet usage")

# ---------------- LOAD MODEL ----------------
@st.cache_resource
def load_model():
    return YOLO("best.pt")   # your trained model path

model = load_model()

# ---------------- TABS ----------------
tab1, tab2 = st.tabs([" Image Detection", " Video Detection"])

# ==================================================
# IMAGE HELMET DETECTION
# ==================================================
with tab1:
    st.header("Image Helmet Detection")

    image_file = st.file_uploader(
        "Upload an Image",
        type=["jpg", "jpeg", "png"]
    )

    if image_file is not None:
        image = Image.open(image_file)
        image_np = np.array(image)

        # YOLO inference
        results = model(image_np, conf=0.4)
        annotated_img = results[0].plot()

        # ---------- DISPLAY IN SMALL WINDOWS ----------
        col1, col2 = st.columns(2)

        with col1:
            st.subheader("Original Image")
            st.image(image, width=350)

        with col2:
            st.subheader("Detection Result")
            st.image(annotated_img, channels="BGR", width=350)

# ==================================================
# VIDEO HELMET DETECTION
# ==================================================
with tab2:
    st.header("Video Helmet Detection")

    video_file = st.file_uploader(
        "Upload a Video",
        type=["mp4", "avi", "mov"]
    )

    if video_file is not None:
        temp_video = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
        temp_video.write(video_file.read())
        temp_video.close()

        cap = cv2.VideoCapture(temp_video.name)
        stframe = st.empty()

        while cap.isOpened():
            ret, frame = cap.read()
            if not ret:
                break

            # YOLO inference
            results = model(frame, conf=0.4)
            annotated_frame = results[0].plot()

            # ---------- SMALL VIDEO WINDOW ----------
            stframe.image(
                annotated_frame,
                channels="BGR",
                width=450
            )

        cap.release()
        os.remove(temp_video.name)

        st.success(" Video processing completed")


st.markdown("""
<br>
<div style='text-align:center; padding:12px; background-color:#111111; border-radius:10px;'>
<span style='color:#AAAAAA; font-size:16px;'>
Designed & Developed by <b style='color:#CCCCCC;'>Yedeedya Injeti</b><br>
Under <b style='color:#B8860B;'>Innomatics Research Labs</b>
</span>
</div>
<br>
""", unsafe_allow_html=True)