Update services/plantation/missing_patch_check.py

services/plantation/missing_patch_check.py

@@ -1,46 +1,96 @@
 import cv2
 import numpy as np
+from ultralytics import YOLO
 import os
 import logging
+from typing import Tuple, List, Dict, Any
 
+# Configure logging
 logging.basicConfig(
     filename="app.log",
     level=logging.INFO,
     format="%(asctime)s - %(levelname)s - %(message)s"
 )
 
-def process_missing_patches(frame):
+# 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/yolov8n.pt"))
+
+# Initialize YOLO model
+try:
+    model = YOLO(MODEL_PATH)
+    logging.info("Loaded YOLOv8n model for missing patch detection.")
+    logging.info(f"Model class names: {model.names}")
+except Exception as e:
+    logging.error(f"Failed to load YOLOv8n model: {str(e)}")
+    model = None
+
+def process_missing_patches(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
+    """
+    Detect missing patches in plantation areas using YOLOv8n.
+
+    Args:
+        frame (np.ndarray): Input frame in BGR format.
+
+    Returns:
+        Tuple[List[Dict[str, Any]], np.ndarray]: A tuple containing:
+            - List of detected missing patches.
+            - Annotated frame with bounding boxes and labels.
+    """
+    # Validate input frame
+    if not isinstance(frame, np.ndarray) or frame.size == 0:
+        logging.error("Invalid input frame provided to missing_patch_check.")
+        return [], frame
+
+    # Check if model is loaded
+    if model is None:
+        logging.error("YOLO model not loaded. Skipping missing patch detection.")
+        return [], frame
+
     try:
-        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
-        _, plant_mask = cv2.threshold(gray, 100, 255, cv2.THRESH_BINARY_INV)
-        contours, _ = cv2.findContours(plant_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+        # Perform YOLO inference
+        results = model(frame, verbose=False)
+        logging.debug("Completed YOLO inference for missing patch detection.")
+    except Exception as e:
+        logging.error(f"Error during YOLO inference: {str(e)}")
+        return [], frame
 
-        detections = []
-        line_counter = 1
+    detections: List[Dict[str, Any]] = []
+    line_counter = 1
 
-        for contour in contours:
-            area = cv2.contourArea(contour)
-            if area < 500:
+    for r in results:
+        for box in r.boxes:
+            conf = float(box.conf[0])
+            if conf < 0.3:
                 continue
-            x, y, w, h = cv2.boundingRect(contour)
-            x_min, y_min, x_max, y_max = x, y, x + w, y + h
+            cls = int(box.cls[0])
+            label = model.names[cls]
+            if label != "missing_patch":
+                continue
+            xyxy = box.xyxy[0].cpu().numpy().astype(int)
+            x_min, y_min, x_max, y_max = xyxy
 
-            detection_label = f"Line {line_counter} - Missing Patch"
+            detection_label = f"Line {line_counter} - Missing Patch (Conf: {conf:.2f})"
             detections.append({
                 "type": "missing_patch",
                 "label": detection_label,
+                "confidence": conf,
                 "coordinates": [x_min, y_min, x_max, y_max]
             })
 
-            color = (139, 69, 19)  # Brown as requested
+            color = (255, 192, 203)  # Pink for missing patches
             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.6, color, 2)
+            cv2.putText(
+                frame,
+                detection_label,
+                (x_min, y_min - 10),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.6,
+                color,
+                2
+            )
             
             line_counter += 1
 
-        logging.info(f"Detected {len(detections)} missing patches in plantation.")
-        return detections, frame
-    except Exception as e:
-        logging.error(f"Error detecting missing patches: {str(e)}")
-        return [], frame
+    logging.info(f"Detected {len(detections)} missing patches in plantation.")
+    return detections, frame