Spaces:
Running
Running
import gradio as gr | |
from ultralytics import YOLO | |
import cv2 | |
import os | |
import numpy as np | |
import sys | |
import asyncio | |
if sys.platform == 'win32': | |
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) | |
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE" | |
# Load your trained model | |
model = YOLO("best.pt") | |
def predict_image(image): | |
# Perform detection | |
detections = model.track(source=image, conf=0.3, iou=0.5, show=False) | |
# Render the detected image | |
detected_image = np.squeeze(detections[0].plot()) | |
# Check if any drones were detected | |
num_drones = len(detections[0]) # Assuming detection results are stored in xywh format | |
message = "Drone detected!" if num_drones > 0 else "No drones detected." | |
return detected_image, message # Return the detected image and the message | |
def predict_video(video_path): | |
# Load the video | |
cap = cv2.VideoCapture(video_path) | |
if not cap.isOpened(): | |
return "Error: Could not open video.", "" | |
drone_detected = False | |
# Prepare to write the output video | |
frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
out_fps = cap.get(cv2.CAP_PROP_FPS) | |
output_path = 'output_video.mp4' | |
out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), out_fps, (frame_width, frame_height)) | |
# Process each frame | |
while cap.isOpened(): | |
ret, frame = cap.read() | |
if not ret: | |
break | |
# Perform detection on the frame | |
results = model.track(frame, imgsz=640, conf=0.3, iou=0.5) | |
# Check if any drones were detected in this frame | |
if len(results[0]) > 0: | |
drone_detected = True | |
# Draw boxes on the frame | |
annotated_frame = np.squeeze(results[0].plot()) | |
# Write the frame to the output video | |
out.write(annotated_frame) | |
# Release everything when done | |
cap.release() | |
out.release() | |
message = "Drone detected in video!" if drone_detected else "No drones detected in video." | |
print("Video processing complete. Saved to:", output_path) | |
return output_path, message # Return the path to the output video and the message | |
with gr.Blocks() as demo: | |
gr.Markdown("### Drone Detection System") | |
with gr.Tab("Introduction"): | |
gr.Markdown("**This Application helps in detection of DRONES in an IMAGE, VIDEO or from your WEBCAM depending on your App mode.**") | |
gr.Markdown("You Don't Necessarily need a Drone to run this app; you can use an image from google.\n\n**SAMPLE OUTPUT:**") | |
gr.Video("Drone Detection.mp4", width=800, height=600) | |
with gr.Tab("Upload Image"): | |
image_input = gr.Image() | |
image_output = gr.Image() | |
alert = gr.Label() | |
image_input.change(fn=predict_image, inputs=image_input, outputs=[image_output, alert]) | |
with gr.Tab("Upload Video"): | |
video_input = gr.Video(sources="upload") | |
video_output = gr.Video(render=True) | |
alert_video = gr.Label() | |
video_input.change(fn=predict_video, inputs=video_input, outputs=[video_output, alert_video]) | |
with gr.Tab("Live"): | |
gr.Markdown("Live detection will be implemented soon.") | |
demo.launch() |