import streamlit as st
from imageai.Detection import ObjectDetection
import os
import time
from PIL import Image

def main():
    st.title("Object Detection App")
    st.write("Upload an image and get object detections!")

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

    if uploaded_file is not None:
        # Display the uploaded image
        image = Image.open(uploaded_file)
        st.image(image, caption='Uploaded Image', use_column_width=True)
        
        # Save the uploaded file temporarily
        with open("temp_image.jpg", "wb") as f:
            f.write(uploaded_file.getbuffer())

        # Button to start detection
        if st.button('Start Detection'):
            # Spinner while model loads
            with st.spinner('Loading model and performing detection...'):
                try:
                    execution_path = os.getcwd()
                    detector = ObjectDetection()
                    detector.setModelTypeAsRetinaNet()
                    detector.setModelPath(os.path.join(execution_path, "retinanet_resnet50_fpn_coco-eeacb38b.pth"))
                    detector.loadModel()

                    # Perform detection
                    detections = detector.detectObjectsFromImage(
                        input_image="temp_image.jpg",
                        output_image_path="output_image.jpg",
                        minimum_percentage_probability=10
                    )

                    # Display results
                    st.success('Detection Complete!')
                    
                    # Display detected image
                    detected_image = Image.open("output_image.jpg")
                    st.image(detected_image, caption='Detected Objects', use_column_width=True)
                    
                    # Display detections with probabilities
                    st.write("### Detected Objects:")
                    for obj in detections:
                        st.write(f"- {obj['name']}: {obj['percentage_probability']:.2f}%")

                except Exception as e:
                    st.error(f"An error occurred: {str(e)}")

            # Clean up temporary files
            if os.path.exists("temp_image.jpg"):
                os.remove("temp_image.jpg")
            if os.path.exists("output_image.jpg"):
                os.remove("output_image.jpg")

if __name__ == "__main__":
    main()