import streamlit as st
from PIL import Image
from ultralytics import YOLO
import tempfile
import os

# Load the YOLO model
model = YOLO("yolov9e.pt")


def process_image(image_path):
    # Run inference
    results = model(image_path)

    # Save the result image
    output_path = os.path.join(tempfile.gettempdir(), "result.jpg")
    results[0].save(filename=output_path)

    return output_path, results[0]


def main():
    st.title("Object Detection App")

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

    if uploaded_file is not None:
        image = Image.open(uploaded_file)

        # Save the uploaded image temporarily
        temp_image_path = os.path.join(tempfile.gettempdir(), uploaded_file.name)
        image.save(temp_image_path)

        if st.button("Process"):
            result_path, result = process_image(temp_image_path)

            st.image(result_path, caption="Detected Objects", use_container_width=True)

            # Display detected objects details
            st.write("### Detected Objects:")

            i = 1
            for box in result.boxes:
                x1, y1, x2, y2 = box.xyxy[0]
                class_id = int(box.cls[0])
                label = model.names[class_id]
                st.write(
                    f"{i}: {label.capitalize()}, **Location:** ({x1:.2f}, {y1:.2f}, {x2:.2f}, {y2:.2f})"
                )
                i += 1


if __name__ == "__main__":
    main()