Spaces:
Sleeping
Sleeping
| 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() |