Update services/under_construction/earthwork_detection.py

services/under_construction/earthwork_detection.py

@@ -2,25 +2,36 @@ import cv2
 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 earthwork detection.")
+except Exception as e:
+    logging.error(f"Failed to load YOLOv8m model: {str(e)}")
+    model = None
 
 def process_earthwork(frame):
-    """
-    Detect earthwork activities in a frame.
-    Args:
-        frame: Input frame (numpy array)
-    Returns:
-        list: List of detected earthwork activities
-        numpy array: Annotated frame with numbered labels
-    """
-    results = model(frame)
+    if model is None:
+        logging.error("YOLO model not loaded. Skipping earthwork detection.")
+        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 +40,11 @@ def process_earthwork(frame):
                 continue
             cls = int(box.cls[0])
             label = model.names[cls]
-            if label != "earthwork":  # Assuming "earthwork" class exists
+            if label != "earthwork":
                 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 +53,12 @@ def process_earthwork(frame):
                 "coordinates": [x_min, y_min, x_max, y_max]
             })
 
-            # Draw bounding box and label
-            color = (255, 228, 181)  # Moccasin for earthwork
+            color = (255, 165, 0)  # Orange 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)} earthwork activities in under_construction.")
     return detections, frame