tomatoseveritysegmentationmodel-l5n39e/app/services/model.py

import numpy as np
import cv2
import time
from typing import Dict, Optional, Any

from app.core.config import settings
from app.utils.image import load_image, resize_image, normalize_image


class TomatoSegmentationModel:
    """
    Model for tomato disease severity segmentation.
    
    This class implements a simple image segmentation model using
    conventional computer vision techniques as a placeholder.
    In a production scenario, this would be replaced with a proper
    deep learning model like U-Net, DeepLabV3, etc.
    """
    
    def __init__(self):
        self.model_name = settings.DEFAULT_MODEL_NAME
        self.model_version = "0.1.0"
        self.severity_classes = settings.SEVERITY_CLASSES
        self.input_size = (224, 224)  # Standard input size
        
    def preprocess_image(self, image_path: str) -> Optional[np.ndarray]:
        """Preprocess an image for analysis."""
        # Load image
        image = load_image(image_path)
        if image is None:
            return None
            
        # Resize
        image = resize_image(image, self.input_size)
        
        # Normalize
        image = normalize_image(image)
        
        return image
        
    def analyze_image(self, image_path: str) -> Dict[str, Any]:
        """
        Analyze a tomato image to determine disease severity.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Dictionary with analysis results including segmentation masks and severity scores
        """
        start_time = time.time()
        
        # Preprocess image
        image = self.preprocess_image(image_path)
        if image is None:
            return {
                "error": "Failed to load or process image",
                "success": False
            }
        
        # Simple color-based segmentation to identify potential disease areas
        # This is a simplified placeholder implementation
        
        # Convert back to BGR for OpenCV
        image_bgr = (image * 255).astype(np.uint8)
        image_bgr = cv2.cvtColor(image_bgr, cv2.COLOR_RGB2BGR)
        
        # Convert to HSV for better color segmentation
        image_hsv = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2HSV)
        
        # Define color ranges for different severity classes
        # These are placeholder values and would need to be calibrated for real use
        color_ranges = {
            "healthy": [(35, 50, 50), (85, 255, 255)],  # Green healthy leaves
            "early_blight": [(15, 50, 50), (35, 255, 255)],  # Yellowish
            "late_blight": [(0, 50, 50), (15, 255, 255)],  # Reddish-brown
            "bacterial_spot": [(0, 0, 0), (180, 50, 100)],  # Dark spots
            "septoria_leaf_spot": [(0, 0, 100), (180, 50, 255)]  # Light spots
        }
        
        # Create segmentation masks and calculate affected areas
        masks = {}
        severity_details = []
        total_pixels = image.shape[0] * image.shape[1]
        
        for severity_class, (lower, upper) in color_ranges.items():
            # Create mask for this color range
            lower = np.array(lower, dtype=np.uint8)
            upper = np.array(upper, dtype=np.uint8)
            mask = cv2.inRange(image_hsv, lower, upper)
            
            # Apply morphological operations to clean up the mask
            kernel = np.ones((5, 5), np.uint8)
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
            
            # Calculate affected area percentage
            affected_pixels = np.count_nonzero(mask)
            affected_percentage = (affected_pixels / total_pixels) * 100
            
            # Generate a random confidence score for demo purposes
            # In a real model, this would be the model's actual confidence
            confidence = min(affected_percentage / 50, 1.0)
            if severity_class == "healthy" and affected_percentage < 10:
                confidence = 0.9  # Boost confidence for low healthy area (likely diseased)
            
            # Store results
            masks[severity_class] = mask.tolist()  # Convert to list for JSON serialization
            severity_details.append({
                "severity_class": severity_class,
                "confidence": float(confidence),
                "affected_area_percentage": float(affected_percentage)
            })
        
        # Determine primary severity class
        severity_details.sort(key=lambda x: x["confidence"], reverse=True)
        primary_severity = severity_details[0]["severity_class"]
        severity_confidence = severity_details[0]["confidence"]
        
        # If healthy has high confidence, it takes precedence
        for detail in severity_details:
            if detail["severity_class"] == "healthy" and detail["confidence"] > 0.7:
                primary_severity = "healthy"
                severity_confidence = detail["confidence"]
                break
        
        # Prepare results
        processing_time = int((time.time() - start_time) * 1000)  # ms
        
        result = {
            "success": True,
            "model_name": self.model_name,
            "model_version": self.model_version,
            "primary_severity": primary_severity,
            "severity_confidence": severity_confidence,
            "severity_details": severity_details,
            "segmentation_data": {
                "image_size": self.input_size,
                "masks": masks
            },
            "processing_time_ms": processing_time
        }
        
        return result


# Singleton instance
segmentation_model = TomatoSegmentationModel()