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

# Load the pre-trained YOLOv8 model
model = YOLO("yolov8x.pt")  # Replace with the path to your model

# Title for the web app
st.title("YOLOv8 Object Detection - Image Upload")

# Instructions
st.write("Upload an image, and YOLOv8 will predict the objects in the image with bounding boxes.")

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

if uploaded_file is not None:
    # Read the uploaded image file and display it
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Convert the image to a numpy array for YOLO processing
    img_array = np.array(image)

    # Make predictions using the model
    results = model.predict(img_array, conf=0.5, iou=0.4)

    # Display the results
    st.write(f"Detected {len(results)} objects.")

    # Annotate the image with bounding boxes
    annotated_img = results[0].plot()

    # Convert the annotated image to a format suitable for Streamlit
    annotated_img_pil = Image.fromarray(annotated_img)

    # Display the annotated image
    st.image(annotated_img_pil, caption="Processed Image with Bounding Boxes", use_column_width=True)