Spaces:
Sleeping
Sleeping
File size: 1,191 Bytes
549a268 fd7e73c 26b4122 fd7e73c 549a268 4b1f1b4 fd7e73c 5a67b15 |
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 |
import cv2
import numpy as np
from ultralytics import YOLO
from typing import List, Tuple, Dict, Any
# Load YOLOv8 model for pothole detection
model = YOLO("models/yolov8n.pt")
def process_potholes(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
"""
Detect potholes in the frame using YOLOv8.
Args:
frame: Input frame as a numpy array.
Returns:
Tuple of (list of detections, annotated frame).
"""
# Perform inference
results = model(frame, classes=[0], conf=0.5) # Class 0 assumed for potholes
detections = []
for i, r in enumerate(results[0].boxes):
x_min, y_min, x_max, y_max = map(int, r.xyxy[0])
conf = float(r.conf)
# Determine severity based on size
area = (x_max - x_min) * (y_max - y_min)
severity = "Severe" if area > 1000 or conf > 0.8 else "Moderate" if area > 500 or conf > 0.6 else "Mild"
detections.append({
"box": [x_min, y_min, x_max, y_max],
"label": f"Pothole {i+1}",
"type": "pothole",
"confidence": conf,
"severity": severity
})
return detections, frame |