services/operations_maintenance/pothole_detection.py

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