Update services/plantation/plant_health.py

services/plantation/plant_health.py

@@ -1,62 +1,70 @@
-# services/plantation/plant_health.py
 import cv2
 import numpy as np
+from ultralytics import YOLO
 import os
 
-def compute_vari(frame):
+# Load YOLOv8n model
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+MODEL_PATH = os.path.abspath(os.path.join(BASE_DIR, "../../models/yolov8n.pt"))
+model = YOLO(MODEL_PATH)
+
+def process_plant_health(frame):
     """
-    Compute VARI (Visible Atmospherically Resistant Index) to assess plant health.
+    Assess plant health by color analysis (simplified as brightness).
     Args:
         frame: Input frame (numpy array)
     Returns:
-        dict: Detection results with numbered labels for unhealthy areas
-        numpy array: Annotated frame
+        list: List of plants with their health status
+        numpy array: Annotated frame with numbered labels
     """
-    # Convert to RGB
-    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-    r, g, b = cv2.split(frame_rgb)
-
-    # Compute VARI: (Green - Red) / (Green + Red - Blue)
-    denominator = (g.astype(float) + r.astype(float) - b.astype(float))
-    denominator[denominator == 0] = 1e-10  # Avoid division by zero
-    vari = (g.astype(float) - r.astype(float)) / denominator
-
-    # Threshold to find unhealthy areas (VARI < 0 indicates poor health)
-    unhealthy_mask = vari < 0
-    contours, _ = cv2.findContours(unhealthy_mask.astype(np.uint8), 
-                                   cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    results = model(frame)
 
     detections = []
     line_counter = 1  # Initialize counter for numbered labels
 
-    # Process unhealthy areas
-    for contour in contours:
-        area = cv2.contourArea(contour)
-        if area < 50:  # Filter small areas
-            continue
-        x, y, w, h = cv2.boundingRect(contour)
-        x_min, y_min, x_max, y_max = x, y, x + w, y + h
-
-        # Add numbered label
-        detection_label = f"Line {line_counter} - Unhealthy Vegetation"
-        detections.append({
-            "label": detection_label,
-            "coordinates": [x_min, y_min, x_max, y_max]
-        })
-
-        # Draw bounding box and label
-        color = (255, 0, 0)  # Red for unhealthy areas
-        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
-
-    return {"detections": detections, "frame": frame}
+    for r in results:
+        for box in r.boxes:
+            conf = float(box.conf[0])
+            if conf < 0.5:
+                continue
+            cls = int(box.cls[0])
+            label = model.names[cls]
+            if label != "plant":  # Assuming "plant" class exists
+                continue
+            xyxy = box.xyxy[0].cpu().numpy()
+            x_min, y_min, x_max, y_max = map(int, xyxy)
 
-def process_plant_health(frame):
-    """
-    Wrapper function for integration with app.py.
-    """
-    result = compute_vari(frame)
-    return result["detections"], result["frame"]
+            # Extract plant region
+            plant_region = frame[y_min:y_max, x_min:x_max]
+            if plant_region.size == 0:
+                continue
+
+            # Convert to HSV and analyze brightness (simplified health check)
+            hsv = cv2.cvtColor(plant_region, cv2.COLOR_BGR2HSV)
+            brightness = hsv[..., 2].mean()
+            if brightness > 150:
+                health = "healthy"
+            elif brightness > 100:
+                health = "moderate"
+            else:
+                health = "unhealthy"
+
+            # Add numbered label
+            detection_label = f"Line {line_counter} - Plant (Health: {health.capitalize()}, Conf: {conf:.2f})"
+            detections.append({
+                "type": "plant",
+                "label": detection_label,
+                "confidence": conf,
+                "coordinates": [x_min, y_min, x_max, y_max],
+                "health": health
+            })
+
+            # Draw bounding box and label
+            color = (34, 139, 34) if health == "healthy" else (255, 165, 0) if health == "moderate" else (255, 0, 0)
+            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
+
+    return detections, frame