import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np
import cv2  # Importing OpenCV (cv2)

# Load the trained YOLO model
model = YOLO("model.pt")  # Replace with the correct path to your model file

def predict(image):
    # Perform prediction on an image
    results_list = model.predict(source=image, conf=0.25)  # Replace 'image.jpg' with your image path
    
    # Iterate over the list of Results objects and process the predictions
    for results in results_list:
        # Plot the results (this will add bounding boxes, labels, etc.)
        annotated_image = results.plot()  # Plotting the results on the image
        
        # Convert the image from BGR to RGB if necessary
        annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
        
        # Convert numpy array to PIL image
        annotated_image = Image.fromarray(annotated_image)
        
        return annotated_image  # Return the image for Gradio to display

# Create Gradio interface
iface = gr.Interface(fn=predict, 
                     inputs=gr.Image(type="pil"),  # Input as a PIL image
                     outputs=gr.Image(type="pil"),  # Output as a PIL image
                     live=True)

# Launch the interface
iface.launch()