Update services/road_safety/pothole_crack_detection.py

services/road_safety/pothole_crack_detection.py

@@ -1,100 +1,40 @@
 import cv2
 import numpy as np
 from ultralytics import YOLO
-import os
-import random
-import logging
-from typing import Tuple, List, Dict, Any
+from typing import List, Tuple, Dict, Any
 
-# Configure logging
-logging.basicConfig(
-    filename="app.log",
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
-)
-
-# Define base directory and model path
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-MODEL_PATH = os.path.abspath(os.path.join(BASE_DIR, "../../models/yolov8m-seg.pt"))
-
-# Initialize YOLO model
-try:
-    model = YOLO(MODEL_PATH)
-    logging.info("Successfully loaded YOLOv8m-seg model for pothole and crack detection.")
-    logging.info(f"Model class names: {model.names}")
-except Exception as e:
-    logging.error(f"Failed to load YOLOv8m-seg model: {str(e)}")
-    model = None
+# Load YOLOv8 model for pothole and crack detection
+model = YOLO("models/yolov8n.pt")  # Use lighter model for faster inference
 
 def detect_potholes_and_cracks(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
-    # Validate input frame
-    if not isinstance(frame, np.ndarray) or frame.size == 0:
-        logging.error("Invalid input frame provided to pothole_crack_detection.")
-        return [], frame
-
-    # Check if model is loaded
-    if model is None:
-        logging.error("YOLO model not loaded. Skipping pothole and crack detection.")
-        return [], frame
-
-    try:
-        # Perform YOLO inference
-        results = model(frame, verbose=False)
-        logging.debug("Completed YOLO inference for pothole and crack detection.")
-    except Exception as e:
-        logging.error(f"Error during YOLO inference: {str(e)}")
-        return [], frame
-
-    detections: List[Dict[str, Any]] = []
-    line_counter = 1
-
-    for result in results:
-        for box in result.boxes:
-            conf = float(box.conf[0])
-            if conf < 0.3:
-                continue
-
-            cls = int(box.cls[0])
-            label = model.names[cls]
-            if label not in ["crack", "pothole", "debris"]:
-                continue
-
-            xyxy = box.xyxy[0].cpu().numpy().astype(int)
-            x_min, y_min, x_max, y_max = xyxy
-
-            severity = None
-            if label == "crack":
-                color = (255, 0, 0)  # Red for cracks
-                severity = random.choice(["low", "medium", "high"])
-            elif label == "pothole":
-                color = (255, 0, 0)  # Red for potholes
-            else:  # debris
-                color = (255, 215, 0)  # Gold for debris
-
-            detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
-            item = {
-                "type": label,
-                "label": detection_label,
-                "confidence": conf,
-                "coordinates": [x_min, y_min, x_max, y_max]
-            }
-            if severity:
-                item["severity"] = severity
-
-            detections.append(item)
-
-            cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
-            cv2.putText(
-                frame,
-                detection_label,
-                (x_min, y_min - 10),
-                cv2.FONT_HERSHEY_SIMPLEX,
-                0.6,
-                color,
-                2
-            )
-
-            line_counter += 1
-
-    logging.info(f"Detected {len(detections)} potholes/cracks/debris in road_safety.")
+    """
+    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)  # Assuming class 0 is pothole, class 1 is 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