nagasurendra's picture
Update app.py
a4bc6c5 verified
import gradio as gr
from ultralytics import YOLO
from PIL import Image
# -----------------------------
# Load YOLO model
# -----------------------------
model = YOLO("./data/best.pt") # make sure this path matches your folder structure
# -----------------------------
# Prediction function
# -----------------------------
def predict(image):
# Run prediction
results = model.predict(image, conf=0.5)
# Annotated image with bounding boxes
result_img = results[0].plot()
# Extract detected labels
detected_labels = results[0].boxes.cls.tolist()
names = results[0].names
detected_objects = [names[int(cls_id)] for cls_id in detected_labels]
# Text output
if detected_objects:
label_text = f"✅ Detected objects: {', '.join(detected_objects)}"
else:
label_text = "❌ No objects detected."
return Image.fromarray(result_img), label_text
# -----------------------------
# Gradio Interface
# -----------------------------
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil"),
outputs=[gr.Image(type="pil", label="Detection Result"), gr.Textbox(label="Detected Objects")],
title="🥤 Bottle Detection with YOLOv11",
description="Upload an image to check if a **bottle** is detected using your trained YOLOv11 model."
)
if __name__ == "__main__":
demo.launch()