Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import cv2 | |
| from ultralytics import YOLO | |
| # Safe load method to handle custom YOLO class during deserialization | |
| def safe_load_yolo_model(path): | |
| # Add necessary safe globals to allow the detection model class during loading | |
| torch.serialization.add_safe_globals([torch, 'ultralytics.nn.tasks.DetectionModel']) | |
| return YOLO(path) | |
| # Load YOLO models | |
| model_yolo11 = safe_load_yolo_model('./data/yolo11n.pt') | |
| model_best = safe_load_yolo_model('./data/best.pt') | |
| def process_video(video): | |
| # Read video input | |
| cap = cv2.VideoCapture(video.name) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| # Create a VideoWriter object to save the output video | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Codec for .mp4 | |
| out = cv2.VideoWriter('output_video.mp4', fourcc, fps, (frame_width, frame_height)) | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Use both YOLO models for detection | |
| results_yolo11 = model_yolo11(frame) | |
| results_best = model_best(frame) | |
| # Combine the results from both models | |
| # For simplicity, we will overlay bounding boxes and labels from both models | |
| for result in results_yolo11: | |
| boxes = result.boxes | |
| for box in boxes: | |
| x1, y1, x2, y2 = map(int, box.xyxy[0].tolist()) | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| label = f"YOLOv11: {box.cls[0]} - {box.conf[0]:.2f}" | |
| cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2) | |
| for result in results_best: | |
| boxes = result.boxes | |
| for box in boxes: | |
| x1, y1, x2, y2 = map(int, box.xyxy[0].tolist()) | |
| cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2) | |
| label = f"Best: {box.cls[0]} - {box.conf[0]:.2f}" | |
| cv2.putText(frame, label, (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2) | |
| # Write the processed frame to the output video | |
| out.write(frame) | |
| cap.release() | |
| out.release() | |
| return 'output_video.mp4' | |
| # Gradio interface | |
| iface = gr.Interface(fn=process_video, inputs=gr.Video(), outputs=gr.Video(), live=True) | |
| # Launch the app | |
| iface.launch() | |