File size: 1,360 Bytes
b291ae2 a4bc6c5 b291ae2 a4bc6c5 b291ae2 a4bc6c5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
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()
|