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

# Load the YOLOv8 model
model = YOLO('yolov8n.pt')

# Streamlit app
def main():
    st.title("Object Detection - General Use")
    st.write("This is a general use object detection space using YOLOv8")
    st.header("Example Result")
    # Display the example result image
    example_image = Image.open("example.jpeg")
    st.image(example_image, caption='Cars are picked out in green and the person is distinguished from the motorcycle', use_column_width=True)
    st.header("Upload an Image")
    # Upload image
    uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
    
    if uploaded_file is not None:
        # Read the uploaded image
        image = Image.open(uploaded_file)
        image_array = np.array(image)
        
        # Perform object detection
        results = model.predict(image_array)
        
        # Display the detected objects
        for result in results:
            boxes = result.boxes
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0]
                x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                cv2.rectangle(image_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
                conf = box.conf[0]
                cls = int(box.cls[0])
                label = f"{model.names[cls]} ({conf:.2f})"
                cv2.putText(image_array, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
        
        # Convert the image array back to PIL Image
        image = Image.fromarray(image_array)
        
        # Display the image with detected objects
        st.image(image, caption='Detected Objects')

if __name__ == "__main__":
    main()