Spaces:
Sleeping
Sleeping
File size: 3,094 Bytes
92ce950 414f345 b763b36 92ce950 b763b36 0e7c662 4dd8158 6bd84a1 0e7c662 4dd8158 92ce950 0e7c662 414f345 0e7c662 e140e6c 414f345 92ce950 0e7c662 414f345 92ce950 b763b36 92ce950 414f345 0e7c662 414f345 0e7c662 414f345 0e7c662 414f345 0e7c662 414f345 0e7c662 414f345 0e7c662 414f345 0e7c662 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 72 73 74 75 |
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/best.pt')
# Class names for YOLO model (replace with actual class names used in your YOLO model)
yolo_classes = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe']
# Class names for best.pt model (assumed classes for crack and pothole)
best_classes = ['Crack', 'Pothole']
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
label = f"YOLO: {yolo_classes[class_id]} - {box.conf[0]:.2f}" # Map class_id to class name
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
label = f"Best: {best_classes[class_id]} - {box.conf[0]:.2f}" # Map class_id to specific labels
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()
|