Update services/operations_maintenance/pothole_detection.py

services/operations_maintenance/pothole_detection.py

@@ -3,15 +3,20 @@ import numpy as np
 from ultralytics import YOLO
 import os
 import logging
+from typing import Tuple, List, 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.pt"))
+
+# Initialize YOLO model
 try:
     model = YOLO(MODEL_PATH)
     logging.info("Loaded YOLOv8m model for pothole detection.")
@@ -19,18 +24,37 @@ except Exception as e:
     logging.error(f"Failed to load YOLOv8m model: {str(e)}")
     model = None
 
-def process_potholes(frame):
+def process_potholes(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
+    """
+    Detect potholes in a video frame using YOLOv8m.
+
+    Args:
+        frame (np.ndarray): Input frame in BGR format.
+
+    Returns:
+        Tuple[List[Dict[str, Any]], np.ndarray]: A tuple containing:
+            - List of detected potholes.
+            - 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_detection.")
+        return [], frame
+
+    # Check if model is loaded
     if model is None:
         logging.error("YOLO model not loaded. Skipping pothole detection.")
         return [], frame
 
     try:
-        results = model(frame)
+        # Perform YOLO inference
+        results = model(frame, verbose=False)
+        logging.debug("Completed YOLO inference for pothole detection.")
     except Exception as e:
         logging.error(f"Error during YOLO inference: {str(e)}")
         return [], frame
 
-    detections = []
+    detections: List[Dict[str, Any]] = []
     line_counter = 1
 
     for r in results:
@@ -55,8 +79,15 @@ def process_potholes(frame):
 
             color = (255, 127, 80)  # Coral (red + orange mix)
             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)
+            cv2.putText(
+                frame,
+                detection_label,
+                (x_min, y_min - 10),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.6,
+                color,
+                2
+            )
             
             line_counter += 1