Update services/plantation/plant_health.py

services/plantation/plant_health.py

@@ -1,77 +1,61 @@
 import cv2
 import numpy as np
-import logging
-from typing import List, Dict, Tuple
+from typing import List, Tuple, Dict, Any
 
-# Setup logging
-logging.basicConfig(
-    filename="app.log",
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
-)
-
-def process_plant_health(frame: np.ndarray) -> Tuple[List[Dict], np.ndarray]:
+def process_plant_health(frame: np.ndarray) -> Tuple[List[Dict[str, Any]], np.ndarray]:
     """
-    Process a frame to assess plant health based on color.
+    Assess the health of plants in the frame based on color.
     Args:
-        frame: Input frame (BGR numpy array)
+        frame: Input frame as a numpy array.
     Returns:
-        Tuple: (List of detection dictionaries with health status, annotated frame)
+        Tuple of (list of detections, annotated frame).
     """
-    try:
-        # Convert frame to HSV
-        hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
-
-        # Define range for green (healthy) and yellow/brown (unhealthy)
-        lower_green = np.array([35, 40, 40])
-        upper_green = np.array([85, 255, 255])
-        green_mask = cv2.inRange(hsv, lower_green, upper_green)
-
-        # Define range for yellow/brown (unhealthy)
-        lower_unhealthy = np.array([20, 40, 40])
-        upper_unhealthy = np.array([35, 255, 255])
-        unhealthy_mask = cv2.inRange(hsv, lower_unhealthy, upper_unhealthy)
-
-        detections = []
-        # Process green regions (healthy plants)
-        contours, _ = cv2.findContours(green_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        for idx, contour in enumerate(contours):
-            area = cv2.contourArea(contour)
-            if area < 300:
-                continue
-            x, y, w, h = cv2.boundingRect(contour)
-            x_min, y_min, x_max, y_max = x, y, x + w, y + h
-            detections.append({
-                "type": "plant",
-                "label": f"Plant {idx + 1} - Healthy",
-                "box": [x_min, y_min, x_max, y_max],
-                "health": "Healthy"
-            })
-            cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
-            cv2.putText(frame, "Healthy", (x_min, y_min - 10),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-
-        # Process unhealthy regions
-        contours, _ = cv2.findContours(unhealthy_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
-        for idx, contour in enumerate(contours):
-            area = cv2.contourArea(contour)
-            if area < 300:
-                continue
-            x, y, w, h = cv2.boundingRect(contour)
-            x_min, y_min, x_max, y_max = x, y, x + w, y + h
-            detections.append({
-                "type": "plant",
-                "label": f"Plant {idx + 1} - Unhealthy",
-                "box": [x_min, y_min, x_max, y_max],
-                "health": "Unhealthy"
-            })
-            cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (255, 0, 0), 2)
-            cv2.putText(frame, "Unhealthy", (x_min, y_min - 10),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 2)
-
-        logging.info(f"Assessed health of {len(detections)} plants in frame.")
-        return detections, frame
-
-    except Exception as e:
-        logging.error(f"Error in plant health assessment: {str(e)}")
-        return [], frame
+    # Convert to HSV color space
+    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
+    
+    # Define range for green (healthy) and yellow/brown (unhealthy)
+    lower_green = np.array([35, 50, 50])
+    upper_green = np.array([85, 255, 255])
+    lower_unhealthy = np.array([20, 50, 50])
+    upper_unhealthy = np.array([35, 255, 255])
+    
+    mask_healthy = cv2.inRange(hsv, lower_green, upper_green)
+    mask_unhealthy = cv2.inRange(hsv, lower_unhealthy, upper_unhealthy)
+    
+    # Combine masks to find plants
+    mask = cv2.bitwise_or(mask_healthy, mask_unhealthy)
+    kernel = np.ones((5, 5), np.uint8)
+    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
+    
+    # Find contours of plants
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    
+    detections = []
+    for i, contour in enumerate(contours):
+        area = cv2.contourArea(contour)
+        if area < 200:  # Ignore small contours
+            continue
+        
+        x, y, w, h = cv2.boundingRect(contour)
+        x_min, y_min, x_max, y_max = x, y, x + w, y + h
+        
+        # Determine health by comparing the area of healthy vs unhealthy pixels
+        roi_hsv = hsv[y_min:y_max, x_min:x_max]
+        healthy_pixels = cv2.countNonZero(cv2.inRange(roi_hsv, lower_green, upper_green))
+        unhealthy_pixels = cv2.countNonZero(cv2.inRange(roi_hsv, lower_unhealthy, upper_unhealthy))
+        total_pixels = healthy_pixels + unhealthy_pixels
+        
+        if total_pixels == 0:
+            continue
+        
+        health_status = "healthy" if healthy_pixels > unhealthy_pixels else "unhealthy"
+        label = f"Plant {i+1} - {health_status.capitalize()}"
+        
+        detections.append({
+            "box": [x_min, y_min, x_max, y_max],
+            "label": label,
+            "type": "plant",
+            "health": health_status
+        })
+    
+    return detections, frame