# Import required libraries
import PIL
import cv2
import streamlit as st
from ultralytics import YOLO
import tempfile


# Replace the relative path to your weight file
model_path = 'https://huggingface.co/spaces/ankitkupadhyay/fire_and_smoke/blob/main/best.pt'

# Setting page layout
st.set_page_config(
    page_title="Object Detection using YOLOv8",  # Setting page title
    page_icon="🤖",     # Setting page icon
    layout="wide",      # Setting layout to wide
    initial_sidebar_state="expanded"    # Expanding sidebar by default
)

# Creating sidebar
with st.sidebar:
    st.header("Image/Video Config")     # Adding header to sidebar
    # Adding file uploader to sidebar for selecting images and videos
    source_file = st.file_uploader(
        "Choose an image or video...", type=("jpg", "jpeg", "png", 'bmp', 'webp', 'mp4'))

    # Model Options
    confidence = float(st.slider(
        "Select Model Confidence", 25, 100, 40)) / 100

# Creating main page heading
st.title("Object Detection using YOLOv8")

# Creating two columns on the main page
col1, col2 = st.columns(2)

# Adding image to the first column if image is uploaded
with col1:
    if source_file:
        # Check if the file is an image
        if source_file.type.split('/')[0] == 'image':
            # Opening the uploaded image
            uploaded_image = PIL.Image.open(source_file)
            # Adding the uploaded image to the page with a caption
            st.image(source_file,
                     caption="Uploaded Image",
                     use_column_width=True
                     )
        else:
            tfile = tempfile.NamedTemporaryFile(delete=False) 
            tfile.write(source_file.read())
            vidcap = cv2.VideoCapture(tfile.name)

try:
    model = YOLO(model_path)
except Exception as ex:
    st.error(
        f"Unable to load model. Check the specified path: {model_path}")
    st.error(ex)

if st.sidebar.button('Detect Objects'):
    if source_file.type.split('/')[0] == 'image':
        res = model.predict(uploaded_image,
                            conf=confidence
                            )
        boxes = res[0].boxes
        res_plotted = res[0].plot()[:, :, ::-1]
        with col2:
            st.image(res_plotted,
                     caption='Detected Image',
                     use_column_width=True
                     )
            try:
                with st.expander("Detection Results"):
                    for box in boxes:
                        st.write(box.xywh)
            except Exception as ex:
                st.write("No image is uploaded yet!")
    else:
        # Open the video file
        success, image = vidcap.read()
        while success:
            res = model.predict(image,
                                conf=confidence
                                )
            boxes = res[0].boxes
            res_plotted = res[0].plot()[:, :, ::-1]
            with col2:
                st.image(res_plotted,
                         caption='Detected Frame',
                         use_column_width=True
                         )
                try:
                    with st.expander("Detection Results"):
                        for box in boxes:
                            st.write(box.xywh)
                except Exception as ex:
                    st.write("No video is uploaded yet!")
            success, image = vidcap.read()