Update services/crack_detection_service.py

services/crack_detection_service.py

@@ -0,0 +1,115 @@
+import cv2
+import numpy as np
+from PIL import Image
+from ultralytics import YOLO
+from collections import defaultdict
+
+# Load YOLOv8 model (yolov8n.pt for lightweight performance)
+try:
+    yolo_model = YOLO('yolov8n.pt')  # Ultralytics will download the compatible version
+except Exception as e:
+    raise RuntimeError(f"Failed to load YOLOv8 model: {str(e)}. Ensure you have an internet connection and the ultralytics package is up to date.")
+
+# Simple tracking state to maintain IDs across frames
+class Tracker:
+    def __init__(self):
+        self.next_id = 0
+        self.objects = {}  # {id: (centroid, type, box, severity, confidence)}
+        self.max_distance = 50  # Max distance to consider same object
+
+    def update(self, detections):
+        new_objects = {}
+        for det in detections:
+            box = det['box']
+            centroid = ((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)
+            
+            best_match_id = None
+            min_dist = float('inf')
+            for obj_id, (obj_centroid, _, _, _, _) in self.objects.items():
+                dist = np.sqrt((centroid[0] - obj_centroid[0])**2 + (centroid[1] - obj_centroid[1])**2)
+                if dist < min_dist and dist < self.max_distance:
+                    min_dist = dist
+                    best_match_id = obj_id
+            
+            if best_match_id is not None:
+                new_objects[best_match_id] = (centroid, det['type'], det['box'], det['severity'], det['confidence'])
+            else:
+                new_objects[self.next_id] = (centroid, det['type'], det['box'], det['severity'], det['confidence'])
+                self.next_id += 1
+        
+        self.objects = new_objects
+        return [(obj_id, obj_data[1], obj_data[2], obj_data[3], obj_data[4]) for obj_id, obj_data in self.objects.items()]
+
+tracker = Tracker()
+
+def segment_road(frame):
+    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+    lower_road = np.array([0, 0, 40])
+    upper_road = np.array([180, 40, 160])
+    mask = cv2.inRange(hsv, lower_road, upper_road)
+    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+    edges = cv2.Canny(gray, 50, 150)
+    edges = cv2.dilate(edges, None, iterations=1)
+    mask = cv2.bitwise_and(mask, mask, mask=cv2.bitwise_not(edges))
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (10, 10))
+    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    if contours:
+        largest_contour = max(contours, key=cv2.contourArea)
+        if cv2.contourArea(largest_contour) > 5000:
+            road_mask = np.zeros_like(mask)
+            cv2.drawContours(road_mask, [largest_contour], -1, 255, -1)
+            return road_mask
+    return np.zeros_like(mask)
+
+def detect_cracks_and_potholes(frame):
+    try:
+        road_mask = segment_road(frame)
+        masked_frame = cv2.bitwise_and(frame, frame, mask=road_mask)
+        gray = cv2.cvtColor(masked_frame, cv2.COLOR_BGR2GRAY)
+        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+        edges = cv2.Canny(blurred, 40, 130)
+        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        
+        cracks = []
+        for contour in contours:
+            if cv2.contourArea(contour) < 150:
+                continue
+            x, y, w, h = cv2.boundingRect(contour)
+            area = w * h
+            if area > 4000:
+                severity = 'Severe'
+            elif area > 1200:
+                severity = 'Moderate'
+            else:
+                severity = 'Minor'
+            aspect_ratio = float(w) / h if h > 0 else 0
+            if 0.3 < aspect_ratio < 7:
+                cracks.append({
+                    'type': 'crack',
+                    'box': [x, y, x + w, y + h],
+                    'severity': severity,
+                    'confidence': 0.95
+                })
+
+        # Skip YOLO detection since yolov8n.pt can't detect potholes
+        potholes = []  # Placeholder; no potholes detected with yolov8n.pt
+
+        filtered_cracks = cracks
+        detections = filtered_cracks + potholes
+        tracked_detections = tracker.update(detections)
+        
+        tracked_items = []
+        for obj_id, obj_type, box, severity, confidence in tracked_detections:
+            tracked_items.append({
+                'id': obj_id,
+                'type': obj_type,
+                'box': box,
+                'severity': severity,
+                'confidence': confidence
+            })
+        
+        return tracked_items
+    except Exception as e:
+        print(f"Error in detection: {str(e)}")
+        return []