File size: 987 Bytes
20ec0fa |
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 |
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) |