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

# Load a pretrained YOLO model
model = YOLO('yolov8n.pt') # Or any other suitable YOLO model

def predict_image_class(input_img):
    """Predicts the class of objects in an image using the YOLO model."""
    results = model.predict(source=input_img, conf=0.25) # Adjust confidence threshold as needed
    
    # Process results and extract classes
    predicted_classes = []
    for r in results:
        boxes = r.boxes
        for box in boxes:
            cls = model.names[int(box.cls)]
            predicted_classes.append(cls)

    return ", ".join(predicted_classes) # Return a string of comma-separated classes

demo = gr.Interface(
    fn=predict_image_class,
    inputs=gr.Image(type="filepath"), # Accept image filepaths as input
    outputs="text",
    title="Image Class Prediction with YOLO",
    description="Upload an image and get predictions of the object classes present."
)

demo.launch(debug=True)