Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| import cv2 | |
| import numpy as np | |
| # Load the trained model | |
| model_path = 'best.pt' # Replace with the path to your trained .pt file | |
| model = YOLO(model_path) | |
| # Function to perform inference on an image | |
| def infer_image(image): | |
| # Convert the image from BGR to RGB | |
| image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # Perform inference | |
| results = model(image_rgb) | |
| # Extract results and annotate image | |
| for result in results: | |
| for box in result.boxes: | |
| x1, y1, x2, y2 = box.xyxy[0] | |
| cls = int(box.cls[0]) | |
| conf = float(box.conf[0]) | |
| # Draw bounding box | |
| cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2) | |
| # Draw label | |
| label = f'{model.names[cls]} {conf:.2f}' | |
| cv2.putText(image, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) | |
| return image | |
| # Create Gradio interface | |
| iface = gr.Interface( | |
| fn=infer_image, | |
| inputs=gr.Image(type="numpy", label="Upload an Image"), | |
| outputs=gr.Image(type="numpy", label="Annotated Image"), | |
| title="YOLOv8 Inference", | |
| description="Upload an image to get object detection results using YOLOv8." | |
| ) | |
| # Launch the app | |
| iface.launch() | |