Update services/overlay_service.py

services/overlay_service.py

@@ -13,15 +13,12 @@ logging.basicConfig(
 def overlay_boxes(frame: np.ndarray, detected_items: List[Dict[str, Any]]) -> np.ndarray:
     """
     Overlay bounding boxes and labels on the frame for detected items.
-
     Args:
-        frame (np.ndarray): Input frame in BGR format.
-        detected_items (List[Dict[str, Any]]): List of detected items with coordinates and labels.
-
+        frame: Input frame in BGR format.
+        detected_items: List of detected items with coordinates and labels.
     Returns:
         np.ndarray: Frame with overlaid bounding boxes and labels.
     """
-    # Validate inputs
     if not isinstance(frame, np.ndarray) or frame.size == 0:
         logging.error("Invalid input frame provided to overlay_service.")
         return frame
@@ -31,35 +28,20 @@ def overlay_boxes(frame: np.ndarray, detected_items: List[Dict[str, Any]]) -> np
 
     try:
         for item in detected_items:
-            # Check for required keys
-            if "coordinates" not in item or "label" not in item:
-                logging.warning(f"Skipping item without coordinates or label: {item}")
+            if "box" not in item or "type" not in item or "severity" not in item:
+                logging.warning(f"Skipping item without box, type, or severity: {item}")
                 continue
 
-            # Extract coordinates and label
-            x_min, y_min, x_max, y_max = item["coordinates"]
-            label = item["label"]
-
-            # Define colors based on detection type
-            type_colors = {
-                "crack": (0, 0, 255),    # Blue for cracks
-                "pothole": (255, 0, 0),  # Red for potholes
-                "signage": (0, 0, 255),  # Blue for signage
-                "default": (0, 255, 255) # Yellow
-            }
-            color = type_colors.get(item.get("type", "default"), type_colors["default"])
-
-            # Draw bounding box and label
+            x_min, y_min, x_max, y_max = map(int, item["box"])
+            severity = item["severity"]
+            if item["type"] == "crack":
+                color = (0, 0, 255) if severity == "Severe" else (0, 255, 255) if severity == "Moderate" else (0, 255, 0) if severity == "Minor" else (255, 165, 0)
+                label = f"Crack: {severity}"
+            else:  # Hole
+                color = (255, 0, 0) if severity == "Severe" else (255, 255, 0) if severity == "Moderate" else (0, 128, 0)
+                label = f"Hole: {severity}"
             cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), color, 2)
-            cv2.putText(
-                frame,
-                label,
-                (x_min, y_min - 10),
-                cv2.FONT_HERSHEY_SIMPLEX,
-                0.6,
-                color,
-                2
-            )
+            cv2.putText(frame, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
 
         logging.info(f"Overlaid {len(detected_items)} items on frame.")
         return frame