Update services/operations_maintenance/crack_detection.py

services/operations_maintenance/crack_detection.py

@@ -1,27 +1,26 @@
-# services/operations_maintenance/crack_detection.py
 import cv2
 import numpy as np
 from ultralytics import YOLO
 import os
+import random
 
-# Load YOLOv8m-seg model
+# Load YOLOv8m-seg model for crack detection
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 MODEL_PATH = os.path.join(BASE_DIR, "../../models/yolov8m-seg.pt")
 model = YOLO(MODEL_PATH)
 
-def detect_cracks(frame):
+def detect_cracks_and_objects(frame):
     """
-    Detect cracks in a frame using YOLOv8m-seg.
+    Detect cracks and other objects in a frame using YOLOv8m-seg.
     Args:
         frame: Input frame (numpy array)
     Returns:
-        dict: Detection results with numbered labels
-        numpy array: Annotated frame
+        list: List of detected items with type, label, coordinates, confidence, and severity
     """
     # Run YOLOv8 inference
     results = model(frame)
 
-    detections = []
+    detected_items = []
     line_counter = 1  # Initialize counter for numbered labels
 
     # Process detections
@@ -32,32 +31,29 @@ def detect_cracks(frame):
                 continue
             cls = int(box.cls[0])
             label = model.names[cls]
-            if label != "crack":  # Assuming "crack" class exists
+            if label not in ["crack", "pothole", "object"]:  # Assuming these classes exist
                 continue
             xyxy = box.xyxy[0].cpu().numpy()
             x_min, y_min, x_max, y_max = map(int, xyxy)
 
+            # Simulate severity for cracks
+            severity = None
+            if label == "crack":
+                severity = random.choice(["low", "medium", "high"])
+            
             # Add numbered label
-            detection_label = f"Line {line_counter} - Crack (Conf: {conf:.2f})"
-            detections.append({
+            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
 
-            # Draw bounding box and label
-            color = (255, 255, 0)  # Yellow for cracks
-            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.5, color, 2)
-            
-            line_counter += 1
+            detected_items.append(item)
 
-    return {"detections": detections, "frame": frame}
+            line_counter += 1
 
-def process_cracks(frame):
-    """
-    Wrapper function for integration with app.py.
-    """
-    result = detect_cracks(frame)
-    return result["detections"], result["frame"]
+    return detected_items