Update services/road_safety/pothole_crack_detection.py

services/road_safety/pothole_crack_detection.py

@@ -6,39 +6,27 @@ import random
 import logging
 from typing import Tuple, List, Dict, Any
 
-# Configure logging to track script execution and errors
+# Configure logging
 logging.basicConfig(
     filename="app.log",
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
-# Define base directory and model path for YOLOv8m-seg
+# 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 with error handling
+# Initialize YOLO model
 try:
     model = YOLO(MODEL_PATH)
     logging.info("Successfully loaded YOLOv8m-seg model for pothole and crack detection.")
-    # Log the model's class names for debugging
     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
 
 def detect_potholes_and_cracks(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
-    """
-    Detect potholes, cracks, and debris in a video frame using YOLOv8m-seg.
-
-    Args:
-        frame (np.ndarray): Input frame in BGR format.
-
-    Returns:
-        Tuple[List[Dict[str, Any]], np.ndarray]: A tuple containing:
-            - List of detected items (dictionaries with type, label, confidence, coordinates, and severity).
-            - Annotated frame with bounding boxes and labels.
-    """
     # Validate input frame
     if not isinstance(frame, np.ndarray) or frame.size == 0:
         logging.error("Invalid input frame provided to pothole_crack_detection.")
@@ -50,50 +38,40 @@ def detect_potholes_and_cracks(frame: np.ndarray) -> Tuple[List[Dict[str, Any]],
         return [], frame
 
     try:
-        # Perform YOLO inference on the frame
-        results = model(frame, verbose=False)  # Suppress verbose output for cleaner logs
+        # 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  # Counter for numbered labels (e.g., "Line 1 - Crack")
+    line_counter = 1
 
-    # Process each detection result
     for result in results:
         for box in result.boxes:
-            # Extract confidence score and ensure it's above threshold
             conf = float(box.conf[0])
-            if conf < 0.3:  # Threshold set to 0.3 for better sensitivity
+            if conf < 0.3:
                 continue
 
-            # Extract class label
             cls = int(box.cls[0])
             label = model.names[cls]
-
-            # Filter for relevant classes
             if label not in ["crack", "pothole", "debris"]:
                 continue
 
-            # Extract bounding box coordinates
             xyxy = box.xyxy[0].cpu().numpy().astype(int)
             x_min, y_min, x_max, y_max = xyxy
 
-            # Assign color based on detection type (as per user requirements)
             severity = None
             if label == "crack":
                 color = (255, 0, 0)  # Red for cracks
                 severity = random.choice(["low", "medium", "high"])
             elif label == "pothole":
-                color = (0, 255, 0)  # Green for potholes
+                color = (255, 0, 0)  # Red for potholes
             else:  # debris
-                color = (0, 0, 255)  # Blue for debris
+                color = (255, 215, 0)  # Gold for debris
 
-            # Create detection label with line number
             detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
-
-            # Create detection dictionary
             item = {
                 "type": label,
                 "label": detection_label,
@@ -105,7 +83,6 @@ def detect_potholes_and_cracks(frame: np.ndarray) -> Tuple[List[Dict[str, Any]],
 
             detections.append(item)
 
-            # Draw bounding box and label on the frame
             cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
             cv2.putText(
                 frame,