Update services/crack_detection_service.py

services/crack_detection_service.py

@@ -1,70 +1,59 @@
 import cv2
 import numpy as np
-
-def detect_cracks_and_holes(frame):
-    try:
-        # Convert frame to grayscale
-        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
-
-        # Edge detection using Canny
-        edges = cv2.Canny(blurred, 30, 100)  # Lower thresholds to detect more cracks
-
-        # Morphological operations to enhance cracks and holes
-        kernel = np.ones((3, 3), np.uint8)
-        dilated = cv2.dilate(edges, kernel, iterations=2)
-        eroded = cv2.erode(dilated, kernel, iterations=1)
-
-        # Find contours
-        contours, _ = cv2.findContours(eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-
-        items = []
-        for contour in contours:
-            # Filter small contours (noise)
-            area = cv2.contourArea(contour)
-            if area < 30:  # Reduced threshold to catch smaller cracks
+from ultralytics import YOLO
+import os
+import random
+
+# 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_and_objects(frame):
+    """
+    Detect cracks and other objects in a frame using YOLOv8m-seg.
+    Args:
+        frame: Input frame (numpy array)
+    Returns:
+        list: List of detected items with type, label, coordinates, confidence, and severity
+    """
+    # Run YOLOv8 inference
+    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])
+            if conf < 0.5:
                 continue
-
-            # Get bounding box
-            x, y, w, h = cv2.boundingRect(contour)
-            area = w * h
-
-            # Calculate aspect ratio and circularity
-            aspect_ratio = float(w) / h if h > 0 else 0
-            perimeter = cv2.arcLength(contour, True)
-            circularity = 4 * np.pi * area / (perimeter * perimeter) if perimeter > 0 else 0
-
-            # Classify as crack or hole
-            if 0.05 < aspect_ratio < 20 and circularity < 0.4:  # Adjusted for more crack detection
-                item_type = 'crack'
-                # Check if it's an underlying crack (longer and thinner)
-                if aspect_ratio > 8:
-                    severity = 'Underlying'
-                elif area > 4000:
-                    severity = 'Severe'
-                elif area > 800:
-                    severity = 'Moderate'
-                else:
-                    severity = 'Minor'
-            elif circularity > 0.6:  # Circular shapes are holes
-                item_type = 'hole'
-                if area > 4000:
-                    severity = 'Severe'
-                elif area > 800:
-                    severity = 'Moderate'
-                else:
-                    severity = 'Minor'
-            else:
-                continue  # Skip objects that don't match crack or hole criteria
-
-            items.append({
-                'type': item_type,
-                'box': [x, y, x + w, y + h],
-                'severity': severity,
-                'confidence': 0.95  # Simulated confidence
-            })
-
-        return items
-    except Exception as e:
-        print(f"Error in detection: {str(e)}")
-        return []
+            cls = int(box.cls[0])
+            label = model.names[cls]
+            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} - {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
+
+            detected_items.append(item)
+
+            line_counter += 1
+
+    return detected_items