Spaces:
Sleeping
Sleeping
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() | |