app.py

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

# Load YOLOv11 Model
model_path = "best.pt"  
model = YOLO(model_path)

def predict(image):
    image = np.array(image)
    results = model(image)

    labels = []
    # Draw bounding boxes and extract labels
    for result in results:
        for box in result.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = box.conf[0]
            cls = int(box.cls[0])
            label = f"{model.names[cls]} {conf:.2f}"
            
            labels.append(label)  # Store detected labels
            
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

    return Image.fromarray(image), labels

# Gradio Interface
iface = gr.Interface(
    fn=predict, 
    inputs="image", 
    outputs=["image", "text"],  # Returning both image and detected labels
    title="YOLOv11 Object Detection"
)

iface.launch()