Spaces:
Sleeping
Sleeping
File size: 2,683 Bytes
92ce950 414f345 b763b36 92ce950 b763b36 0e7c662 4dd8158 6bd84a1 0e7c662 4dd8158 cb86f50 92ce950 414f345 0e7c662 e140e6c 414f345 92ce950 0e7c662 414f345 92ce950 b763b36 92ce950 414f345 0e7c662 414f345 0e7c662 414f345 0e7c662 414f345 0e7c662 156c9ac 26080f3 414f345 0e7c662 414f345 0e7c662 414f345 0e7c662 156c9ac 26080f3 414f345 0e7c662 414f345 92ce950 289a4db b763b36 414f345 b763b36 f5710ab 414f345 92ce950 414f345 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import gradio as gr
import torch
import cv2
from ultralytics import YOLO
# Load YOLO models
def safe_load_yolo_model(path):
torch.serialization.add_safe_globals([torch, 'ultralytics.nn.tasks.DetectionModel'])
return YOLO(path)
# Load the models
model_yolo11 = safe_load_yolo_model('./data/yolo11n.pt')
model_best = safe_load_yolo_model('./data/best2.pt')
def process_video(video):
# Open the video using OpenCV
cap = cv2.VideoCapture(video)
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 VideoWriter to save 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
# Detect with YOLOv11 (general object detection model)
results_yolo11 = model_yolo11(frame)
# Detect with best.pt (specialized model for cracks and potholes)
results_best = model_best(frame)
# Draw bounding boxes and labels for YOLOv11 (General Object Detection)
for result in results_yolo11:
boxes = result.boxes
for box in boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
class_id = int(box.cls[0]) # Class index for YOLO
# Use model's built-in class names
label = f"{model_yolo11.names[class_id]} - {box.conf[0]:.2f}"
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
# Draw bounding boxes and labels for best.pt (Crack and Pothole detection)
for result in results_best:
boxes = result.boxes
for box in boxes:
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
class_id = int(box.cls[0]) # Class index for best.pt
# Use model's built-in class names for best.pt
label = f"{model_best.names[class_id]} - {box.conf[0]:.2f}"
cv2.rectangle(frame, (x1, y1), (x2, y2), (255, 0, 0), 2)
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()
|