Update services/crack_detection_service.py

services/crack_detection_service.py

@@ -9,21 +9,22 @@ 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_and_objects(frame):
+import random
+
+def detect_cracks(frame, model):
     """
-    Detect cracks and other objects in a frame using YOLOv8m-seg.
+    Detect cracks in a frame using YOLOv8.
     Args:
         frame: Input frame (numpy array)
+        model: YOLO model
     Returns:
-        list: List of detected items with type, label, coordinates, confidence, and severity
+        list: List of detected cracks with type, label, coordinates, confidence, and severity
     """
-    # Run YOLOv8 inference
+    # Run YOLOv8 inference for cracks
     results = model(frame)
-
     detected_items = []
     line_counter = 1  # Initialize counter for numbered labels
 
-    # Process detections
     for r in results:
         for box in r.boxes:
             conf = float(box.conf[0])
@@ -31,29 +32,127 @@ def detect_cracks_and_objects(frame):
                 continue
             cls = int(box.cls[0])
             label = model.names[cls]
-            if label not in ["crack", "pothole", "object"]:  # Assuming these classes exist
+            if label != "crack":  # Process only cracks
                 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"])
+            severity = random.choice(["low", "medium", "high"])
             
             # Add numbered label
-            detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
+            detection_label = f"Line {line_counter} - Crack (Conf: {conf:.2f})"
+            item = {
+                "type": label,
+                "label": detection_label,
+                "confidence": conf,
+                "coordinates": [x_min, y_min, x_max, y_max],
+                "severity": severity
+            }
+
+            detected_items.append(item)
+            line_counter += 1
+
+    return detected_items
+
+def detect_potholes(frame, model):
+    """
+    Detect potholes in a frame using YOLOv8.
+    Args:
+        frame: Input frame (numpy array)
+        model: YOLO model
+    Returns:
+        list: List of detected potholes with type, label, coordinates, and confidence
+    """
+    # Run YOLOv8 inference for potholes
+    results = model(frame)
+    detected_items = []
+    line_counter = 1  # Initialize counter for numbered labels
+
+    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 != "pothole":  # Process only potholes
+                continue
+            xyxy = box.xyxy[0].cpu().numpy()
+            x_min, y_min, x_max, y_max = map(int, xyxy)
+
+            # Add numbered label
+            detection_label = f"Line {line_counter} - Pothole (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
 
             detected_items.append(item)
+            line_counter += 1
+
+    return detected_items
+
+def detect_objects(frame, model):
+    """
+    Detect objects in a frame using YOLOv8.
+    Args:
+        frame: Input frame (numpy array)
+        model: YOLO model
+    Returns:
+        list: List of detected objects with type, label, coordinates, and confidence
+    """
+    # Run YOLOv8 inference for other objects
+    results = model(frame)
+    detected_items = []
+    line_counter = 1  # Initialize counter for numbered labels
+
+    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 != "object":  # Process only objects
+                continue
+            xyxy = box.xyxy[0].cpu().numpy()
+            x_min, y_min, x_max, y_max = map(int, xyxy)
+
+            # Add numbered label
+            detection_label = f"Line {line_counter} - Object (Conf: {conf:.2f})"
+            item = {
+                "type": label,
+                "label": detection_label,
+                "confidence": conf,
+                "coordinates": [x_min, y_min, x_max, y_max]
+            }
 
+            detected_items.append(item)
             line_counter += 1
 
-    return detected_items
+    return detected_items
+
+def detect_items_in_sequence(frame, model):
+    """
+    Run crack, pothole, and object detection sequentially.
+    Args:
+        frame: Input frame (numpy array)
+        model: YOLO model
+    Returns:
+        list: List of detected items (crack, pothole, object)
+    """
+    detected_items = []
+
+    # Detect cracks first
+    detected_items.extend(detect_cracks(frame, model))
+
+    # Detect potholes second
+    detected_items.extend(detect_potholes(frame, model))
+
+    # Detect objects third
+    detected_items.extend(detect_objects(frame, model))
+
+    return detected_items