Update services/road_safety/barrier_check.py

services/road_safety/barrier_check.py

@@ -1,26 +1,38 @@
 import cv2
 import numpy as np
+import numpy as np
 from ultralytics import YOLO
 import os
+import logging
+
+logging.basicConfig(
+    filename="app.log",
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s"
+)
 
-# Load YOLOv8m model
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 MODEL_PATH = os.path.abspath(os.path.join(BASE_DIR, "../../models/yolov8m.pt"))
-model = YOLO(MODEL_PATH)
+try:
+    model = YOLO(MODEL_PATH)
+    logging.info("Loaded YOLOv8m model for barrier check.")
+except Exception as e:
+    logging.error(f"Failed to load YOLOv8m model: {str(e)}")
+    model = None
 
 def process_barriers(frame):
-    """
-    Detect road barriers in a frame.
-    Args:
-        frame: Input frame (numpy array)
-    Returns:
-        list: List of detected barriers
-        numpy array: Annotated frame with numbered labels
-    """
-    results = model(frame)
+    if model is None:
+        logging.error("YOLO model not loaded. Skipping barrier check.")
+        return [], frame
+
+    try:
+        results = model(frame)
+    except Exception as e:
+        logging.error(f"Error during YOLO inference: {str(e)}")
+        return [], frame
 
     detections = []
-    line_counter = 1  # Initialize counter for numbered labels
+    line_counter = 1
 
     for r in results:
         for box in r.boxes:
@@ -29,12 +41,11 @@ def process_barriers(frame):
                 continue
             cls = int(box.cls[0])
             label = model.names[cls]
-            if label != "barrier":  # Assuming "barrier" class exists
+            if label != "barrier":
                 continue
-            xyxy = box.xyxy[0].cpu().numpy()
-            x_min, y_min, x_max, y_max = map(int, xyxy)
+            xyxy = box.xyxy[0].cpu().numpy().astype(int)
+            x_min, y_min, x_max, y_max = xyxy
 
-            # Add numbered label
             detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
             detections.append({
                 "type": label,
@@ -43,12 +54,12 @@ def process_barriers(frame):
                 "coordinates": [x_min, y_min, x_max, y_max]
             })
 
-            # Draw bounding box and label
-            color = (255, 215, 0)  # Gold for barriers
+            color = (0, 0, 255)  # Blue as requested
             cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
             cv2.putText(frame, detection_label, (x_min, y_min - 10), 
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
             
             line_counter += 1
 
+    logging.info(f"Detected {len(detections)} barriers in road_safety.")
     return detections, frame