Update services/under_construction/earthwork_detection.py

services/under_construction/earthwork_detection.py

@@ -1,4 +1,3 @@
-# services/under_construction/earthwork_detection.py
 import cv2
 import numpy as np
 from ultralytics import YOLO
@@ -6,58 +5,38 @@ import os
 
 # Load YOLOv8m model
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-MODEL_PATH = os.path.join(BASE_DIR, "../../models/yolov8m.pt")
+MODEL_PATH = os.path.abspath(os.path.join(BASE_DIR, "../../models/yolov8m.pt"))
 model = YOLO(MODEL_PATH)
 
-def detect_earthwork(frame):
+def process_earthwork(frame):
     """
-    Detect earthwork-related objects in a frame using YOLOv8m.
+    Detect earthwork activities in a frame.
     Args:
         frame: Input frame (numpy array)
     Returns:
-        dict: Detection results with numbered labels
-        numpy array: Annotated frame
+        list: List of detected earthwork activities
+        numpy array: Original frame (overlays handled elsewhere)
     """
-    # Run YOLOv8 inference
     results = model(frame)
 
     detections = []
-    line_counter = 1  # Initialize counter for numbered labels
 
-    # Process detections
     for r in results:
         for box in r.boxes:
             conf = float(box.conf[0])
-            if conf < 0.5:  # Confidence threshold
+            if conf < 0.5:
                 continue
             cls = int(box.cls[0])
-            label = model.names[cls]  # Get class name (e.g., "truck", "person")
-            if label not in ["truck", "person", "excavator", "bulldozer"]:
+            label = model.names[cls]
+            if label != "earthwork":  # Assuming "earthwork" class exists
                 continue
-            xyxy = box.xyxy[0].cpu().numpy()  # Bounding box coordinates
+            xyxy = box.xyxy[0].cpu().numpy()
             x_min, y_min, x_max, y_max = map(int, xyxy)
 
-            # Add numbered label
-            detection_label = f"Line {line_counter} - {label.capitalize()} (Conf: {conf:.2f})"
             detections.append({
-                "label": detection_label,
+                "type": label,
                 "confidence": conf,
                 "coordinates": [x_min, y_min, x_max, y_max]
             })
 
-            # Draw bounding box and label on frame
-            color = (0, 255, 0)  # Green for earthwork objects
-            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)
-            
-            line_counter += 1  # Increment counter
-
-    return {"detections": detections, "frame": frame}
-
-def process_earthwork(frame):
-    """
-    Wrapper function for integration with app.py.
-    """
-    result = detect_earthwork(frame)
-    return result["detections"], result["frame"]
+    return detections, frame