Create pothole_crack_detection.py

services/road_safety/pothole_crack_detection.py

@@ -0,0 +1,77 @@
+import cv2
+import numpy as np
+from ultralytics import YOLO
+import os
+import random
+import logging
+
+logging.basicConfig(
+    filename="app.log",
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
+
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+MODEL_PATH = os.path.abspath(os.path.join(BASE_DIR, "../../models/yolov8m-seg.pt"))
+try:
+    model = YOLO(MODEL_PATH)
+    logging.info("Loaded YOLOv8m-seg model for pothole and crack detection.")
+except Exception as e:
+    logging.error(f"Failed to load YOLOv8m-seg model: {str(e)}")
+    model = None
+
+def detect_potholes_and_cracks(frame):
+    if model is None:
+        logging.error("YOLO model not loaded. Skipping pothole and crack detection.")
+        return [], frame
+
+    try:
+        results = model(frame)
+    except Exception as e:
+        logging.error(f"Error during YOLO inference: {str(e)}")
+        return [], frame
+
+    detections = []
+    line_counter = 1
+
+    for r in results:
+        for box in r.boxes:
+            conf = float(box.conf[0])
+            if conf < 0.5:
+                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":
+                severity = (255, 0, 0)  # Red
+                severity = random.choice(["low", "medium", "high"])
+            elif label == "pothole":
+                color = (0, 255, 0)  # Green
+            else:  # debris
+                color = (0, 0, 255)  # Blue
+
+            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_SIMPLEXSIMPLEX, 0.6, color, 2)
+
+            line_counter += 1
+
+    logging.info(f"Detected {len(detections)} potholes/cracks in road_safety.")
+    return detections, frame