Update services/operations_maintenance/crack_detection.py

services/operations_maintenance/crack_detection.py

@@ -1,76 +1,54 @@
 import cv2
 import numpy as np
-from ultralytics import YOLO
-import os
-import random
 import logging
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Tuple
+from ultralytics import YOLO
 
-# Configure logging
+# Setup logging
 logging.basicConfig(
     filename="app.log",
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
-# Define base directory and model path
-BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-MODEL_PATH = os.path.abspath(os.path.join(BASE_DIR, "../../models/yolov8m.pt"))
-
-# Initialize YOLO model
-try:
-    model = YOLO(MODEL_PATH)
-    logging.info("Loaded YOLOv8m model for crack detection.")
-    logging.info(f"Model class names: {model.names}")
-except Exception as e:
-    logging.error(f"Failed to load YOLOv8m model: {str(e)}")
-    model = None
-
-def detect_cracks_and_objects(frame: np.ndarray) -> List[Dict[str, Any]]:
-    # Validate input frame
-    if not isinstance(frame, np.ndarray) or frame.size == 0:
-        logging.error("Invalid input frame provided to crack_detection.")
-        return []
-
-    # Check if model is loaded
-    if model is None:
-        logging.error("YOLO model not loaded. Skipping crack detection.")
-        return []
-
+def detect_cracks_and_holes(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
+    """
+    Detect cracks and holes in a frame using YOLO.
+    Args:
+        frame: Input frame (numpy array)
+    Returns:
+        tuple: List of detected items (dicts with type, box, severity, etc.) and the input frame
+    """
     try:
-        # Perform YOLO inference
-        results = model(frame, verbose=False)
-        logging.debug("Completed YOLO inference for crack detection.")
+        # 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 during YOLO inference: {str(e)}")
-        return []
-
-    detections: List[Dict[str, Any]] = []
-    line_counter = 1
-
-    for r in results:
-        for box in r.boxes:
-            conf = float(box.conf[0])
-            if conf < 0.3:
-                continue
-            cls = int(box.cls[0])
-            label = model.names[cls]
-            if label != "crack":
-                continue
-            xyxy = box.xyxy[0].cpu().numpy().astype(int)
-            x_min, y_min, x_max, y_max = xyxy
-
-            severity = random.choice(["low", "medium", "high"])
-            detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
-            detections.append({
-                "type": label,
-                "label": detection_label,
-                "confidence": conf,
-                "coordinates": [x_min, y_min, x_max, y_max],
-                "severity": severity
-            })
-
-            line_counter += 1
-
-    logging.info(f"Detected {len(detections)} cracks in operations_maintenance.")
-    return detections
+        logging.error(f"Error detecting cracks/holes: {str(e)}")
+        return [], frame