Update services/operations_maintenance/crack_detection.py

services/operations_maintenance/crack_detection.py

@@ -3,57 +3,83 @@ import numpy as np
 from ultralytics import YOLO
 import os
 import random
+import logging
+from typing import List, Dict, Any
 
-# Load YOLOv8m-seg model for crack detection
+# Configure 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-seg.pt"))
-model = YOLO(MODEL_PATH)
+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):
+def detect_cracks_and_objects(frame: np.ndarray) -> List[Dict[str, Any]]:
     """
-    Detect cracks and other objects in a frame using YOLOv8m-seg.
+    Detect cracks and other objects in a video frame using YOLOv8m.
+
     Args:
-        frame: Input frame (numpy array)
+        frame (np.ndarray): Input frame in BGR format.
+
     Returns:
-        list: List of detected items with type, label, coordinates, confidence, and severity
+        List[Dict[str, Any]]: List of detected items with type, label, confidence, coordinates, and severity.
     """
-    # Run YOLOv8 inference
-    results = model(frame)
+    # Validate input frame
+    if not isinstance(frame, np.ndarray) or frame.size == 0:
+        logging.error("Invalid input frame provided to crack_detection.")
+        return []
 
-    detected_items = []
-    line_counter = 1  # Initialize counter for numbered labels
+    # Check if model is loaded
+    if model is None:
+        logging.error("YOLO model not loaded. Skipping crack detection.")
+        return []
+
+    try:
+        # Perform YOLO inference
+        results = model(frame, verbose=False)
+        logging.debug("Completed YOLO inference for crack detection.")
+    except Exception as e:
+        logging.error(f"Error during YOLO inference: {str(e)}")
+        return []
+
+    detections: List[Dict[str, Any]] = []
+    line_counter = 1
 
-    # Process detections
     for r in results:
         for box in r.boxes:
             conf = float(box.conf[0])
-            if conf < 0.5:
+            if conf < 0.3:
                 continue
             cls = int(box.cls[0])
             label = model.names[cls]
-            if label not in ["crack", "pothole", "object"]:  # Assuming these classes exist
+            if label != "crack":
                 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
+            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})"
-            item = {
+            detections.append({
                 "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)
+                "coordinates": [x_min, y_min, x_max, y_max],
+                "severity": severity
+            })
 
             line_counter += 1
 
-    return detected_items
+    logging.info(f"Detected {len(detections)} cracks in operations_maintenance.")
+    return detections