Update services/crack_detection_service.py

services/crack_detection_service.py

@@ -1,55 +1,68 @@
-import importlib.util
-import sys
-from transformers import DetrImageProcessor, DetrForObjectDetection
-import torch
-from PIL import Image
 import cv2
-import random
+import numpy as np
 
-# Ensure timm is installed
-if not importlib.util.find_spec("timm"):
+def detect_cracks_and_holes(frame):
     try:
-        import pip
-        pip.main(['install', 'timm==0.9.2'])
-    except Exception as e:
-        raise RuntimeError("Failed to install timm. Please ensure it is included in requirements.txt and restart the runtime.")
+        # Convert frame to grayscale
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
 
-# Load model with error handling
-try:
-    processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50", force_download=True)
-    model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50", force_download=True)
-except Exception as e:
-    raise RuntimeError(f"Failed to load DETR model: {str(e)}")
+        # Edge detection using Canny
+        edges = cv2.Canny(blurred, 50, 150)
 
-def detect_cracks_and_objects(frame):
-    try:
-        image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-        inputs = processor(images=image, return_tensors="pt")
-        outputs = model(**inputs)
+        # Morphological operations to enhance cracks and holes
+        kernel = np.ones((3, 3), np.uint8)
+        dilated = cv2.dilate(edges, kernel, iterations=1)
+        eroded = cv2.erode(dilated, kernel, iterations=1)
 
-        target_sizes = torch.tensor([image.size[::-1]])
-        results = processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[0]
+        # Find contours
+        contours, _ = cv2.findContours(eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
 
         items = []
-        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
-            if score >= 0.9:
-                box_area = (box[2] - box[0]) * (box[3] - box[1])
-                label_str = model.config.id2label[label.item()]
-                if "crack" in label_str.lower() or random.random() > 0.3:  # Simulate crack detection
-                    severity = "Large" if box_area > 1000 else "Small"
-                    items.append({
-                        'type': 'crack',
-                        'box': box.tolist(),
-                        'severity': severity,
-                        'confidence': score.item()
-                    })
+        for contour in contours:
+            # Filter small contours (noise)
+            area = cv2.contourArea(contour)
+            if area < 50:
+                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.1 < aspect_ratio < 10 and circularity < 0.5:  # Long, thin shapes are cracks
+                item_type = 'crack'
+                # Check if it's an underlying crack (longer and thinner)
+                if aspect_ratio > 5:
+                    severity = 'Underlying'
+                elif area > 5000:
+                    severity = 'Severe'
+                elif area > 1000:
+                    severity = 'Moderate'
+                else:
+                    severity = 'Minor'
+            elif circularity > 0.7:  # Circular shapes are holes
+                item_type = 'hole'
+                if area > 5000:
+                    severity = 'Severe'
+                elif area > 1000:
+                    severity = 'Moderate'
                 else:
-                    items.append({
-                        'type': 'object',
-                        'box': box.tolist(),
-                        'label': label_str,
-                        'confidence': score.item()
-                    })
+                    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: