Detectify / app.py
arpit13's picture
Update app.py
9bacbe9 verified
raw
history blame
941 Bytes
import gradio as gr
import cv2
import numpy as np
from ultralytics import YOLO # Import YOLO from ultralytics
# Load the YOLOv5 model
model = YOLO('yolov5s') # Use the YOLOv5s pre-trained model
# Function to run inference on an image
def run_inference(image):
# Convert the image from PIL format to a format compatible with OpenCV
image = np.array(image)
# Run YOLOv5 inference
results = model.predict(source=image, save=False, conf=0.25, stream=False)
# Annotate the image with detected objects
annotated_image = results[0].plot() # Use YOLO's built-in plotting function
return annotated_image
# Create the Gradio interface
interface = gr.Interface(
fn=run_inference,
inputs=gr.Image(type="pil"),
outputs=gr.Image(type="numpy"),
title="YOLOv5 Object Detection",
description="Upload an image to run YOLOv5 object detection and see the results."
)
# Launch the app
interface.launch()