Spaces:
Build error
Build error
File size: 2,924 Bytes
8085eab | 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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | import gradio as gr
import torch
from ultralytics import YOLO
import numpy as np
from PIL import Image
import os
# Load YOLOv8 Small model
print("Loading YOLOv8 Small model...")
model = YOLO('yolov8s.pt')
print("Model loaded successfully!")
def predict(image, confidence_threshold, iou_threshold):
"""
Run YOLOv8 inference on the input image
"""
try:
# Run inference
results = model(
image,
conf=confidence_threshold,
iou=iou_threshold
)
# Get the first result (single image)
result = results[0]
# Create annotated image
annotated_image = result.plot()
# Convert BGR to RGB for proper display
annotated_image = Image.fromarray(annotated_image[..., ::-1])
# Get detection info
detections = []
if result.boxes is not None:
for box in result.boxes:
class_id = int(box.cls[0])
class_name = model.names[class_id]
confidence = float(box.conf[0])
detections.append(f"{class_name}: {confidence:.2f}")
detection_text = "\n".join(detections) if detections else "No objects detected"
return annotated_image, detection_text
except Exception as e:
return None, f"Error: {str(e)}"
# Create Gradio interface
with gr.Blocks(title="YOLOv8 Object Detection") as demo:
gr.Markdown("# 🔍 YOLOv8 Small Object Detection")
gr.Markdown("Upload an image to detect objects using YOLOv8s model")
with gr.Row():
with gr.Column():
input_image = gr.Image(type="pil", label="Upload Image")
confidence_slider = gr.Slider(
minimum=0.1,
maximum=1.0,
value=0.5,
step=0.05,
label="Confidence Threshold"
)
iou_slider = gr.Slider(
minimum=0.1,
maximum=1.0,
value=0.5,
step=0.05,
label="IoU Threshold"
)
submit_btn = gr.Button("Detect Objects", variant="primary")
with gr.Column():
output_image = gr.Image(type="pil", label="Detection Results")
output_text = gr.Textbox(label="Detected Objects", lines=10)
# Connect the function
submit_btn.click(
fn=predict,
inputs=[input_image, confidence_slider, iou_slider],
outputs=[output_image, output_text]
)
# Add examples
gr.Examples(
examples=[
["https://ultralytics.com/images/bus.jpg", 0.5, 0.5],
],
inputs=[input_image, confidence_slider, iou_slider],
outputs=[output_image, output_text],
fn=predict,
cache_examples=False
)
if __name__ == "__main__":
demo.launch() |