File size: 1,274 Bytes
6a88941 | 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 | import gradio as gr
from ultralytics import YOLO
from PIL import Image
import numpy as np
import cv2 # Importing OpenCV (cv2)
# Load the trained YOLO model
model = YOLO("model.pt") # Replace with the correct path to your model file
def predict(image):
# Perform prediction on an image
results_list = model.predict(source=image, conf=0.25) # Replace 'image.jpg' with your image path
# Iterate over the list of Results objects and process the predictions
for results in results_list:
# Plot the results (this will add bounding boxes, labels, etc.)
annotated_image = results.plot() # Plotting the results on the image
# Convert the image from BGR to RGB if necessary
annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
# Convert numpy array to PIL image
annotated_image = Image.fromarray(annotated_image)
return annotated_image # Return the image for Gradio to display
# Create Gradio interface
iface = gr.Interface(fn=predict,
inputs=gr.Image(type="pil"), # Input as a PIL image
outputs=gr.Image(type="pil"), # Output as a PIL image
live=True)
# Launch the interface
iface.launch()
|