gnithin/drone_detection

1# prompt: create a gradio app to use the trained yolo model where we can pass an image or video as input and gradio should display the output2 3import gradio as gr4from ultralytics import YOLO5import cv26import numpy as np7 8# Load the trained YOLO model9model = YOLO('best.pt')10 11def predict_image(image):12    results = model.predict(source=image, save=False)13    14    # Process the results and draw bounding boxes15    im = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)16    for data in results[0].boxes.data:17          xmin, ymin, xmax, ymax = int(data[0]), int(data[1]), int(data[2]), int(data[3])18          cv2.rectangle(im,(xmin,ymin),(xmax,ymax),(0,0,255),2)19          cv2.putText(im,20                      results[0].names[int(data[-1])],21                      (xmin, ymin),22                      cv2.FONT_HERSHEY_SIMPLEX,23                      .5, (255, 0, 0), 1)24    25    im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)26    return im27 28 29def predict_video(video):30    results = model.predict(source=video, save=False) # Assuming save=False is correct31    # Placeholder for video processing (not implemented yet)32    return "Video processing not yet implemented"33 34 35iface = gr.Interface(36    fn=predict_image, 37    inputs=gr.Image(),38    outputs=gr.Image(type="numpy"),39    title="Drone Detection",40    description="Upload an image or video to detect drones."41)42 43 44iface.launch()45