import streamlit as st
from ultralytics import YOLO
import os
from os import listdir, remove
from PIL import Image
import cv2


import requests

def download_file(url, destination_folder):
    """
    Download a file from a URL and save it to a local destination.
    
    Args:
        url (str): The URL of the file to download.
        destination_folder (str): The local folder where the file should be saved.
    """
    # Create the destination folder if it doesn't exist
    if not os.path.exists(destination_folder):
        os.makedirs(destination_folder)
    
    # Extract filename from URL
    filename = url.split('/')[-1]
    destination_path = os.path.join(destination_folder, filename)
    
    # Download the file
    with open(destination_path, 'wb') as f:
        response = requests.get(url)
        f.write(response.content)
    
    return destination_path

url = 'https://yolo-v8-models.s3.ap-south-1.amazonaws.com/modelv2/best.pt'  
destination_folder = 'models'  
    
downloaded_file = download_file(url, destination_folder)

# Load the model
model = YOLO('models/best.pt')

# Streamlit UI
st.title("Object Detection with Ultralytics")

# Upload image or video
uploaded_file = st.file_uploader("Choose an image or video", type=["jpg", "jpeg", "png","mp4"])

# Demo section
st.header("Demo")
col1, col2 = st.columns(2)

# Display image in first column
with col1:
    st.image("demoimg.jpeg", caption="Annotated Image", use_column_width=True)

  # Display video in second column with adjusted width
with col2:
    st.write(f'<div style="width: {300};height :500">', unsafe_allow_html=True)
    st.video("demovideo.mp4")
    st.write('</div>', unsafe_allow_html=True)

if uploaded_file is not None:
    # Check if the uploaded file is a video
    if uploaded_file.type.startswith("video/"):
        # Progress bar to show the progress of object detection
        progress_bar = st.progress(0)
        
        st.header(uploaded_file.name)

        # Perform object detection
        with st.spinner('Performing object detection...'):
            for percent_complete in range(100):
                result = model.predict(source=uploaded_file, conf=0.2, save=True ,stream=True)
                for i in result : i.save("video.mp4")
                progress_bar.progress(percent_complete + 1)
        
        st.success(f"Video saved successfully ")
        # Perform object detection
    

    else:
        
        # Read the uploaded image
        image = Image.open(uploaded_file)
        
        img_name = "converted_image.jpg"
        image.save(img_name)

        # Perform object detection
        result = model.predict(source=img_name, conf=0.2, save=True)

        # Save the output image
        img_save_path = "output/"
        os.makedirs(img_save_path, exist_ok=True)
        for r in result:
            r.save(filename=os.path.join(img_save_path, uploaded_file.name))

        st.success("Detected Object")
        # Display the output image
        st.image(os.path.join(img_save_path, uploaded_file.name), caption="Detected Objects", use_column_width=True)