import gradio as gr
import torch
from ultralytics import YOLO
import cv2
import numpy as np

# Load the trained YOLO model
model = YOLO("yolo11n.pt")  # Ensure this matches the model file in your repo

# Define the function for making predictions
def predict(image):
    # Convert to OpenCV format
    img = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    # Run YOLOv8 inference
    results = model(img)

    detected_objects = []  # Store detected object names and confidence

    # Draw bounding boxes on the image
    for r in results:
        for box in r.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])  # Get bounding box coordinates
            conf = float(box.conf[0])  # Confidence score
            cls = int(box.cls[0])  # Class index
            
            label = f"{model.names[cls]} {conf:.2f}"
            detected_objects.append(label)  # Save object name and confidence
            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    # Convert back to PIL format
    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB), "\n".join(detected_objects) if detected_objects else "No objects detected."

# Create the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil"), gr.Textbox(label="Detected Objects")],
    title="YOLOv8 Fruit Detection",
    description="Upload an image and the model will detect objects and list them."
)

# Launch the Gradio app
if __name__ == "__main__":
    iface.launch()