Spaces:
Runtime error
Runtime error
File size: 1,348 Bytes
edb10ef 609231f edb10ef 609231f 8ec8a00 edb10ef a592e90 609231f 8ec8a00 609231f edb10ef |
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 |
import cv2
import numpy as np
from ultralytics import YOLO
from typing import List, Tuple, Dict, Any
# Load YOLOv8 model for pothole and crack detection
model = YOLO("models/yolov8n.pt")
def detect_potholes_and_cracks(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
"""
Detect potholes and cracks 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, 1], conf=0.5) # Class 0: pothole, Class 1: crack
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)
cls = int(r.cls)
dtype = "hole" if cls == 0 else "crack"
label = f"{dtype.capitalize()} {i+1}"
# Determine severity based on confidence and 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": label,
"type": dtype,
"confidence": conf,
"severity": severity
})
return detections, frame |