Update services/operations_maintenance/crack_detection.py

services/operations_maintenance/crack_detection.py

@@ -1,54 +1,58 @@
 import cv2
 import numpy as np
-import logging
-from typing import List, Dict, Any, Tuple
-from ultralytics import YOLO
-
-# Setup logging
-logging.basicConfig(
-    filename="app.log",
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
-)
+from typing import List, Tuple, Dict, Any
 
 def detect_cracks_and_holes(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
     """
-    Detect cracks and holes in a frame using YOLO.
+    Detect cracks and holes in the frame using edge detection and contour analysis.
     Args:
-        frame: Input frame (numpy array)
+        frame: Input frame as a numpy array.
     Returns:
-        tuple: List of detected items (dicts with type, box, severity, etc.) and the input frame
+        Tuple of (list of detections, annotated frame).
     """
-    try:
-        # Load YOLO model
-        model = YOLO("models/yolov8m.pt")
-        # Run inference
-        results = model(frame)
-
-        detections: List[Dict[str, Any]] = []
-        line_counter = 1
-
-        for result in results:
-            boxes = result.boxes.xyxy.cpu().numpy()
-            confidences = result.boxes.conf.cpu().numpy()
-            classes = result.boxes.cls.cpu().numpy()
-
-            for box, conf, cls in zip(boxes, confidences, classes):
-                x_min, y_min, x_max, y_max = map(int, box)
-                type_ = "crack" if cls == 0 else "hole"  # Assume class 0=crack, 1=hole
-                severity = "Severe" if conf > 0.8 else "Moderate" if conf > 0.5 else "Minor"
-                label = f"Line {line_counter} - {type_.capitalize()} ({severity})"
-                detections.append({
-                    "type": type_,
-                    "label": label,
-                    "box": [x_min, y_min, x_max, y_max],
-                    "severity": severity,
-                    "confidence": float(conf)
-                })
-                line_counter += 1
-
-        logging.info(f"Detected {len(detections)} cracks/holes in frame.")
-        return detections, frame
-    except Exception as e:
-        logging.error(f"Error detecting cracks/holes: {str(e)}")
-        return [], frame
+    # Convert to grayscale
+    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+    
+    # Apply Gaussian blur to reduce noise
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    
+    # Edge detection using Canny
+    edges = cv2.Canny(blurred, 50, 150)
+    
+    # Dilate edges to connect nearby edges
+    kernel = np.ones((3, 3), np.uint8)
+    dilated = cv2.dilate(edges, kernel, iterations=1)
+    
+    # Find contours
+    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    
+    detections = []
+    for i, contour in enumerate(contours):
+        # Calculate the area of the contour
+        area = cv2.contourArea(contour)
+        if area < 100:  # Ignore small contours
+            continue
+        
+        # Get bounding box
+        x, y, w, h = cv2.boundingRect(contour)
+        x_min, y_min, x_max, y_max = x, y, x + w, y + h
+        
+        # Determine if it's a crack or hole based on shape and area
+        perimeter = cv2.arcLength(contour, True)
+        circularity = 4 * np.pi * area / (perimeter * perimeter) if perimeter > 0 else 0
+        
+        # Classify as hole if more circular, crack if elongated
+        dtype = "hole" if circularity > 0.5 else "crack"
+        label = f"{dtype.capitalize()} {i+1}"
+        
+        # Determine severity based on area
+        severity = "Severe" if area > 1000 else "Moderate" if area > 500 else "Mild"
+        
+        detections.append({
+            "box": [x_min, y_min, x_max, y_max],
+            "label": label,
+            "type": dtype,
+            "severity": severity
+        })
+    
+    return detections, frame