demeter

  1from ultralytics import YOLO
  2import cv2
  3import json
  4import os
  5import logging
  6import numpy as np
  7import base64
  8from io import BytesIO
  9from PIL import Image
 10
 11# Setup basic logging
 12logging.basicConfig(level=logging.INFO)
 13logger = logging.getLogger(__name__)
 14
 15class VisionAgent:
 16    def __init__(self, model_path=None):
 17        logger.info("👁️ Initializing Vision Agent (Doctor)...")
 18        
 19        # 1. Find the project root
 20        if model_path:
 21            default_model = model_path
 22        else:
 23            current_file = os.path.abspath(__file__)
 24            agent_dir = os.path.dirname(os.path.dirname(current_file)) # agent/
 25            default_model = os.path.join(agent_dir, "model", "plant_disease_model.pt")
 26        
 27        self.model_name = default_model
 28        
 29        # 2. Load Model with Fallback
 30        try:
 31            if os.path.exists(self.model_name):
 32                logger.info(f"✅ Found plant disease model at: {self.model_name}")
 33                self.model = YOLO(self.model_name)
 34            else:
 35                logger.warning(f"⚠️ Custom model not found. Using generic YOLOv8n.")
 36                self.model = YOLO("yolov8n.pt")
 37                self.model_name = "yolov8n.pt"
 38            
 39            # Optimization
 40            self.model.to('cpu')
 41            
 42        except Exception as e:
 43            logger.error(f"❌ Critical Error loading model: {e}")
 44            self.model = None
 45
 46    def analyze_frame(self, image_b64):
 47        """
 48        Scans a base64 encoded image for pests, diseases, or growth stages.
 49        """
 50        if not self.model:
 51            return {"error": "Model not initialized"}
 52            
 53        if not image_b64:
 54            return {"error": "No image data provided"}
 55
 56        try:
 57            # 3. Decode Base64 to Image
 58            # Handle data URI scheme if present (e.g., "data:image/png;base64,...")
 59            if "," in image_b64:
 60                image_b64 = image_b64.split(",")[1]
 61            
 62            image_data = base64.b64decode(image_b64)
 63            image = Image.open(BytesIO(image_data))
 64
 65            # 4. Run Inference
 66            # YOLO can accept PIL Images directly
 67            results = self.model.predict(image, conf=0.25, save=False, verbose=False)
 68            result = results[0]
 69            
 70            detections = []
 71            summary_counts = {}
 72            
 73            for box in result.boxes:
 74                class_id = int(box.cls[0])
 75                label = self.model.names[class_id]
 76                confidence = float(box.conf[0])
 77                
 78                detections.append({
 79                    "object": label,
 80                    "confidence": round(confidence, 2),
 81                    "box": [round(x, 2) for x in box.xywhn[0].tolist()] 
 82                })
 83                summary_counts[label] = summary_counts.get(label, 0) + 1
 84
 85            # 5. Health Logic
 86            health_status = "HEALTHY"
 87            visual_alert = False
 88            
 89            if not detections:
 90                # If generic model, it might just see nothing. 
 91                # If disease model, empty usually means healthy.
 92                if "yolov8n" in self.model_name:
 93                    health_status = "NO_OBJECTS_DETECTED"
 94                else:
 95                    health_status = "HEALTHY"
 96            else:
 97                for label in summary_counts:
 98                    label_lower = label.lower()
 99                    # Keywords that imply sickness
100                    sick_keywords = ['spot', 'rot', 'blight', 'mildew', 'rust', 'virus', 'miner', 'mite', 'wilt']
101                    if any(x in label_lower for x in sick_keywords) and "healthy" not in label_lower:
102                        health_status = "DISEASE_DETECTED"
103                        visual_alert = True
104                        break
105
106            return {
107                "status": "Success",
108                "model_used": self.model_name,
109                "health_assessment": health_status,
110                "visual_alert": visual_alert,
111                "object_counts": summary_counts,
112                "detailed_detections": detections
113            }
114            
115        except Exception as e:
116           
117            logger.error(f"Error during analysis: {e}")
118            return {"error": str(e)}