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

# ----------- Page config ----------------
st.set_page_config(
    page_title='Fire Detection',
    page_icon='',
    layout='wide'
)

st.title("Fire Detection")
st.write('Upload an image or video to detect Fire')

# -------------- load model detection -----------
@st.cache_resource
def load_model():
    return YOLO('best.pt')

model = load_model()

# switching tabs
x, y = st.tabs(['Image Detection', 'Video Detection'])

# ==================================================
# Image Detection
# ==================================================
with x:
    st.header('Image Detection')

    img_path = st.file_uploader('Please upload an image')

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

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

        # Convert BGR to RGB
        annot_img = cv2.cvtColor(annot_img, cv2.COLOR_BGR2RGB)
        annot_img = cv2.cvtColor(annot_img, cv2.COLOR_BGR2RGB)


        # Display side by side
        ori_img, pre_img = st.columns(2)

        with ori_img:
            st.markdown('#### ***Original Image***')
            st.image(image, width=400)

        with pre_img:
            st.markdown('#### ***Detected Image***')
            st.image(annot_img, width=400)

# ==================================================
# Video Detection
# ==================================================
with y:
    st.header("Video Fire 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)

        col1, col2 = st.columns(2)

        with col1:
            st.markdown("#### **Original Video**")
            orig_frame = st.empty()

        with col2:
            st.markdown("#### **Detected Video**")
            pred_frame = st.empty()

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

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

            orig_frame.image(frame, channels="BGR", width=400)
            pred_frame.image(annotated_frame, channels="BGR", width=400)

        cap.release()
        os.remove(temp_video.name)
        st.success("Video processing completed")

# ---------------- Sample Test Images & Videos ----------------
st.markdown("---")
st.subheader("Try with Sample Images / Videos")
st.write("Don't have files? Use the samples below to test the model.")

SAMPLE_IMAGES = {
    "Fire Image 1": "pcb1.jpg",
    "Fire Image 2": "pcb4.jpg",
    "Fire Image 3": "pcb5.jpg"
}

SAMPLE_VIDEOS = {
    "Fire Video 1": "v1.mp4",
    "Fire Video 2": "v2.mp4",
    "Fire Video 3": "v3.mp4",
    "Fire Video 4": "v4.mp4"
}


col1, col2 = st.columns(2)

# -------- Sample Images --------
with col1:
    st.markdown("###  Sample Images")
    selected_img = st.selectbox(
        "Choose a sample image",
        ["None"] + list(SAMPLE_IMAGES.keys())
    )

    if selected_img != "None":
        img_path = SAMPLE_IMAGES[selected_img]
        image = Image.open(img_path)
        st.image(image, caption=selected_img, use_container_width=True)

        if st.button("Detect Fire in Image"):
            results = model(image)
            annotated_img = results[0].plot()
            st.image(annotated_img, caption="Detection Result", use_container_width=True)

# -------- Sample Videos --------
with col2:
    st.markdown("### 🎥 Sample Videos")
    selected_vid = st.selectbox(
        "Choose a sample video",
        ["None"] + list(SAMPLE_VIDEOS.keys())
    )

    if selected_vid != "None":
        video_path = SAMPLE_VIDEOS[selected_vid]
        st.video(video_path)

        if st.button("Detect Fire in Video"):
            cap = cv2.VideoCapture(video_path)
            stframe = st.empty()

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

                results = model(frame)
                annotated_frame = results[0].plot()
                stframe.image(annotated_frame, channels="BGR", use_container_width=True)

            cap.release()





# ---------------- Footer ----------------
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)