src/app.py

import cv2
import numpy as np
from matplotlib import pyplot as plt

class FruitDiseaseDetector:
    def __init__(self):
        self.disease_mask = None
        self.healthy_mask = None
        self.processed_image = None
        
    def remove_background(self, image):
        """
        Remove background from fruit image using multiple segmentation techniques
        Returns the fruit mask and the image with background removed
        """
        # Method 1: GrabCut algorithm (most effective for fruits)
        fruit_mask_grabcut = self._grabcut_segmentation(image)
        
        # Method 2: Color-based segmentation
        fruit_mask_color = self._color_based_segmentation(image)
        
        # Method 3: Edge-based segmentation
        fruit_mask_edge = self._edge_based_segmentation(image)
        
        # Combine all methods using voting
        combined_mask = self._combine_masks([fruit_mask_grabcut, fruit_mask_color, fruit_mask_edge])
        
        # Post-process the mask
        final_mask = self._post_process_mask(combined_mask)
        
        # Apply mask to image
        result_image = image.copy()
        result_image[final_mask == 0] = [0, 0, 0]  # Set background to black
        
        return final_mask, result_image
    
    def _grabcut_segmentation(self, image):
        """Use GrabCut algorithm for mango foreground/background separation"""
        height, width = image.shape[:2]
        
        # Initialize mask
        mask = np.zeros((height, width), np.uint8)
        
        # Define rectangle around the mango (mangoes are typically oval/elongated)
        # Adjust margins for mango shape - less margin on sides, more on top/bottom
        margin_x = int(width * 0.10)   # Reduced horizontal margin for mango width
        margin_y = int(height * 0.12)  # Slightly more vertical margin for mango length
        rect = (margin_x, margin_y, width - 2*margin_x, height - 2*margin_y)
        
        # Initialize background and foreground models
        bgd_model = np.zeros((1, 65), np.float64)
        fgd_model = np.zeros((1, 65), np.float64)
        
        # Apply GrabCut with more iterations for better mango segmentation
        cv2.grabCut(image, mask, rect, bgd_model, fgd_model, 8, cv2.GC_INIT_WITH_RECT)
        
        # Create binary mask (0 for background, 1 for foreground)
        fruit_mask = np.where((mask == 2) | (mask == 0), 0, 255).astype(np.uint8)
        
        return fruit_mask
    
    def _color_based_segmentation(self, image):
        """Segment mango using mango-specific color characteristics"""
        # Convert to HSV for better color segmentation
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        
        # Mango-specific color ranges (calibrated for mango dataset)
        # Green mangoes (unripe) - broader green range for mangoes
        lower_green1 = np.array([35, 25, 25])
        upper_green1 = np.array([85, 255, 255])
        
        # Yellow/Orange mangoes (ripe) - extended range for mango ripeness
        lower_yellow = np.array([10, 30, 30])
        upper_yellow = np.array([35, 255, 255])
        
        # Red/Orange mangoes (very ripe) - specific to mango varieties
        lower_red1 = np.array([0, 30, 30])
        upper_red1 = np.array([15, 255, 255])
        lower_red2 = np.array([165, 30, 30])
        upper_red2 = np.array([180, 255, 255])
        
        # Yellowish-green mangoes (semi-ripe)
        lower_yellow_green = np.array([25, 20, 40])
        upper_yellow_green = np.array([45, 255, 255])
        
        # Create masks for mango color ranges
        mask_green = cv2.inRange(hsv, lower_green1, upper_green1)
        mask_yellow = cv2.inRange(hsv, lower_yellow, upper_yellow)
        mask_red1 = cv2.inRange(hsv, lower_red1, upper_red1)
        mask_red2 = cv2.inRange(hsv, lower_red2, upper_red2)
        mask_yellow_green = cv2.inRange(hsv, lower_yellow_green, upper_yellow_green)
        
        # Combine all mango color masks
        fruit_mask = cv2.bitwise_or(mask_green, mask_yellow)
        fruit_mask = cv2.bitwise_or(fruit_mask, mask_red1)
        fruit_mask = cv2.bitwise_or(fruit_mask, mask_red2)
        fruit_mask = cv2.bitwise_or(fruit_mask, mask_yellow_green)
        
        # Include areas with moderate saturation and brightness (mango skin variations)
        saturation = hsv[:, :, 1]  # Saturation channel
        value = hsv[:, :, 2]       # Value channel
        
        # Mango skin can have lower saturation but still be fruit
        _, moderate_sat_mask = cv2.threshold(saturation, 30, 255, cv2.THRESH_BINARY)
        _, bright_mask = cv2.threshold(value, 50, 255, cv2.THRESH_BINARY)
        
        # Additional mask for pale/light mango areas
        mango_skin_mask = cv2.bitwise_and(moderate_sat_mask, bright_mask)
        
        # Final combination optimized for mangoes
        fruit_mask = cv2.bitwise_or(fruit_mask, mango_skin_mask)
        
        return fruit_mask
    
    def _edge_based_segmentation(self, image):
        """Use edge detection and contour analysis for segmentation"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # Apply Gaussian blur
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        # Edge detection
        edges = cv2.Canny(blurred, 50, 150)
        
        # Find contours
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # Create mask
        mask = np.zeros(gray.shape, np.uint8)
        
        if contours:
            # Find the largest contour (assuming it's the fruit)
            largest_contour = max(contours, key=cv2.contourArea)
            
            # Only consider if the contour is reasonably large
            if cv2.contourArea(largest_contour) > (gray.shape[0] * gray.shape[1] * 0.1):
                cv2.fillPoly(mask, [largest_contour], 255)
        
        return mask
    
    def _combine_masks(self, masks):
        """Combine multiple masks using majority voting"""
        height, width = masks[0].shape
        combined = np.zeros((height, width), dtype=np.uint8)
        
        # Convert masks to binary (0 or 1)
        binary_masks = [(mask > 127).astype(np.uint8) for mask in masks]
        
        # Sum all masks
        mask_sum = np.sum(binary_masks, axis=0)
        
        # Use majority voting (at least 2 out of 3 methods agree)
        combined[mask_sum >= 2] = 255
        
        return combined
    
    def _post_process_mask(self, mask):
        """Clean up the mask using morphological operations"""
        # Remove small noise
        kernel_small = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        cleaned = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel_small)
        
        # Fill small holes
        kernel_large = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (15, 15))
        cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, kernel_large)
        
        # Find the largest connected component (should be the main fruit)
        num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(cleaned, connectivity=8)
        
        if num_labels > 1:
            # Find largest component (excluding background)
            largest_label = 1 + np.argmax(stats[1:, cv2.CC_STAT_AREA])
            
            # Create mask with only the largest component
            final_mask = np.zeros_like(cleaned)
            final_mask[labels == largest_label] = 255
        else:
            final_mask = cleaned
        
        return final_mask

    def preprocess_image(self, image):
        """Preprocess the input image for disease detection"""
        # Convert to different color spaces for better analysis
        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
        
        # Apply Gaussian blur to reduce noise
        blurred = cv2.GaussianBlur(image, (5, 5), 0)
        
        return blurred, hsv, lab
    
    def detect_diseased_areas(self, image, fruit_mask=None):
        """
        Detect mango diseases using color-based segmentation and texture analysis
        Calibrated for: Alternaria, Anthracnose, Aspergillus (Black Mould), Lasiodiplodia (Stem Rot)
        """
        blurred, hsv, lab = self.preprocess_image(image)
        
        # Method 1: Detect Anthracnose (dark circular spots with orange/pink halos)
        disease_mask1 = self._detect_anthracnose(hsv, lab)
        
        # Method 2: Detect Alternaria (brown/black irregular spots)
        disease_mask2 = self._detect_alternaria(hsv, lab)
        
        # Method 3: Detect Aspergillus (Black Mould Rot - dark, fuzzy patches)
        disease_mask3 = self._detect_aspergillus(hsv, lab)
        
        # Method 4: Detect Lasiodiplodia (Stem and Rot - soft, dark areas)
        disease_mask4 = self._detect_lasiodiplodia(hsv, lab)
        
        # Method 5: General texture-based detection for rough/irregular surfaces
        disease_mask5 = self._detect_texture_anomalies(blurred)
        
        # Method 6: Edge-based detection for irregular boundaries
        disease_mask6 = self._detect_irregular_edges(blurred)
        
        # Combine disease-specific methods first (higher confidence)
        primary_disease_mask = cv2.bitwise_or(disease_mask1, disease_mask2)
        primary_disease_mask = cv2.bitwise_or(primary_disease_mask, disease_mask3)
        primary_disease_mask = cv2.bitwise_or(primary_disease_mask, disease_mask4)
        
        # Secondary detection (texture and edges) - use only if there's significant evidence
        secondary_mask = cv2.bitwise_and(disease_mask5, disease_mask6)
        
        # Combine primary and secondary with different weights
        combined_mask = cv2.bitwise_or(primary_disease_mask, secondary_mask)
        
        # Apply fruit mask to limit detection to mango area only
        if fruit_mask is not None:
            combined_mask = cv2.bitwise_and(combined_mask, fruit_mask)
        
        # Post-processing: More aggressive noise removal for better accuracy
        kernel_small = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
        kernel_medium = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
        
        # Remove small noise (more aggressive)
        combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel_small)
        combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_OPEN, kernel_medium)
        
        # Fill small holes but not too aggressively
        kernel_close = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
        combined_mask = cv2.morphologyEx(combined_mask, cv2.MORPH_CLOSE, kernel_close)
        
        # Filter out very small regions (likely noise)
        contours, _ = cv2.findContours(combined_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        filtered_mask = np.zeros_like(combined_mask)
        min_area = 75  # Balanced minimum area for good detection
        
        for contour in contours:
            if cv2.contourArea(contour) >= min_area:
                cv2.fillPoly(filtered_mask, [contour], 255)
        
        return filtered_mask
    
    def _detect_anthracnose(self, hsv, lab):
        """Detect Anthracnose disease - dark circular spots with orange/pink halos"""
        # Anthracnose characteristics: Dark centers with lighter halos
        # Use both HSV and LAB for better detection
        
        # Dark spots in HSV (less conservative thresholds)
        lower_dark = np.array([0, 40, 0])
        upper_dark = np.array([180, 255, 80])
        dark_mask = cv2.inRange(hsv, lower_dark, upper_dark)
        
        # Reddish-brown areas (anthracnose lesions) - broader range
        lower_reddish = np.array([0, 50, 15])
        upper_reddish = np.array([25, 255, 140])
        reddish_mask = cv2.inRange(hsv, lower_reddish, upper_reddish)
        
        # LAB color space for better brown/dark detection
        l_channel = lab[:, :, 0]
        a_channel = lab[:, :, 1]
        
        # Less conservative dark areas with reddish tint
        _, dark_lab = cv2.threshold(l_channel, 85, 255, cv2.THRESH_BINARY_INV)
        _, red_lab = cv2.threshold(a_channel, 135, 255, cv2.THRESH_BINARY)
        
        lab_anthracnose = cv2.bitwise_and(dark_lab, red_lab)
        
        # Combine HSV and LAB detections - use OR instead of AND for better sensitivity
        anthracnose_mask = cv2.bitwise_or(dark_mask, reddish_mask)
        anthracnose_mask = cv2.bitwise_or(anthracnose_mask, lab_anthracnose)
        
        return anthracnose_mask
    
    def _detect_alternaria(self, hsv, lab):
        """Detect Alternaria disease - brown/black irregular spots with concentric rings"""
        # Alternaria has characteristic brown to black lesions
        
        # Less conservative brown to black spots in HSV
        lower_brown1 = np.array([5, 60, 5])
        upper_brown1 = np.array([25, 255, 100])
        
        lower_brown2 = np.array([0, 40, 3])
        upper_brown2 = np.array([20, 255, 80])
        
        brown_mask1 = cv2.inRange(hsv, lower_brown1, upper_brown1)
        brown_mask2 = cv2.inRange(hsv, lower_brown2, upper_brown2)
        
        # Very dark areas (advanced alternaria) - less conservative
        lower_black = np.array([0, 30, 0])
        upper_black = np.array([180, 255, 50])
        black_mask = cv2.inRange(hsv, lower_black, upper_black)
        
        # LAB space detection for brown lesions - less strict
        l_channel = lab[:, :, 0]
        _, dark_lab = cv2.threshold(l_channel, 70, 255, cv2.THRESH_BINARY_INV)
        
        # Combine alternaria indicators - use OR for better sensitivity
        alternaria_mask = cv2.bitwise_or(brown_mask1, brown_mask2)
        alternaria_mask = cv2.bitwise_or(alternaria_mask, black_mask)
        alternaria_mask = cv2.bitwise_or(alternaria_mask, dark_lab)
        
        return alternaria_mask
    
    def _detect_aspergillus(self, hsv, lab):
        """Detect Aspergillus (Black Mould Rot) - dark, fuzzy patches with irregular borders"""
        # Aspergillus appears as very dark, irregular patches with possible greenish tints
        
        # Very dark areas in HSV - less conservative
        lower_black = np.array([0, 20, 0])
        upper_black = np.array([180, 255, 40])
        black_mask = cv2.inRange(hsv, lower_black, upper_black)
        
        # Dark greenish areas (aspergillus can have greenish tint) - broader range
        lower_dark_green = np.array([40, 40, 3])
        upper_dark_green = np.array([80, 255, 60])
        dark_green_mask = cv2.inRange(hsv, lower_dark_green, upper_dark_green)
        
        # Dark bluish-green areas (specific to aspergillus) - broader range
        lower_blue_green = np.array([85, 30, 5])
        upper_blue_green = np.array([120, 255, 80])
        blue_green_mask = cv2.inRange(hsv, lower_blue_green, upper_blue_green)
        
        # LAB space for very dark areas - less conservative
        l_channel = lab[:, :, 0]
        b_channel = lab[:, :, 2]
        
        _, very_dark_lab = cv2.threshold(l_channel, 50, 255, cv2.THRESH_BINARY_INV)
        
        # Blue-green tint in LAB space (aspergillus characteristic) - broader range
        _, blue_tint = cv2.threshold(b_channel, 110, 255, cv2.THRESH_BINARY_INV)
        
        lab_aspergillus = cv2.bitwise_and(very_dark_lab, blue_tint)
        
        # Combine aspergillus indicators - use OR for better sensitivity
        aspergillus_mask = cv2.bitwise_or(black_mask, very_dark_lab)
        aspergillus_mask = cv2.bitwise_or(aspergillus_mask, dark_green_mask)
        aspergillus_mask = cv2.bitwise_or(aspergillus_mask, blue_green_mask)
        aspergillus_mask = cv2.bitwise_or(aspergillus_mask, lab_aspergillus)
        
        return aspergillus_mask
    
    def _detect_lasiodiplodia(self, hsv, lab):
        """Detect Lasiodiplodia (Stem and Rot disease) - soft, dark areas with brown/black coloration"""
        # Lasiodiplodia characteristics: dark, soft areas with brown/black coloration
        
        # Brown rot areas - less conservative
        lower_rot = np.array([8, 50, 10])
        upper_rot = np.array([30, 255, 120])
        rot_mask = cv2.inRange(hsv, lower_rot, upper_rot)
        
        # Dark spots with higher saturation (active rot) - less strict
        lower_active_rot = np.array([0, 60, 5])
        upper_active_rot = np.array([20, 255, 100])
        active_rot_mask = cv2.inRange(hsv, lower_active_rot, upper_active_rot)
        
        # Very dark areas (advanced lasiodiplodia) - broader range
        lower_very_dark = np.array([0, 30, 0])
        upper_very_dark = np.array([30, 255, 60])
        very_dark_mask = cv2.inRange(hsv, lower_very_dark, upper_very_dark)
        
        # LAB space for detecting rot areas - less conservative
        l_channel = lab[:, :, 0]
        a_channel = lab[:, :, 1]
        
        _, dark_rot = cv2.threshold(l_channel, 70, 255, cv2.THRESH_BINARY_INV)
        _, red_tint = cv2.threshold(a_channel, 130, 255, cv2.THRESH_BINARY)
        
        lab_rot = cv2.bitwise_and(dark_rot, red_tint)
        
        # Combine lasiodiplodia indicators - use OR for better sensitivity
        lasiodiplodia_mask = cv2.bitwise_or(rot_mask, active_rot_mask)
        lasiodiplodia_mask = cv2.bitwise_or(lasiodiplodia_mask, very_dark_mask)
        lasiodiplodia_mask = cv2.bitwise_or(lasiodiplodia_mask, lab_rot)
        
        return lasiodiplodia_mask
    
    def _detect_black_mould_rot(self, hsv, lab):
        """Detect Black Mould Rot - dark, fuzzy patches with irregular borders (legacy method)"""
        # This method is now replaced by _detect_aspergillus but kept for compatibility
        return self._detect_aspergillus(hsv, lab)
    
    def _detect_stem_rot(self, hsv, lab):
        """Detect Stem and Rot disease - soft, dark areas typically near stem end (legacy method)"""
        # This method is now replaced by _detect_lasiodiplodia but kept for compatibility
        return self._detect_lasiodiplodia(hsv, lab)
    
    def _detect_brown_spots(self, hsv):
        """Detect brown/dark spots typical of fruit diseases (legacy method)"""
        # Define HSV range for brown/dark diseased areas
        # Brown spots: Low saturation, low value, wide hue range
        lower_brown = np.array([0, 20, 20])
        upper_brown = np.array([30, 255, 120])
        
        # Create mask for brown areas
        brown_mask1 = cv2.inRange(hsv, lower_brown, upper_brown)
        
        # Additional brown range (reddish-brown)
        lower_brown2 = np.array([150, 20, 20])
        upper_brown2 = np.array([180, 255, 120])
        brown_mask2 = cv2.inRange(hsv, lower_brown2, upper_brown2)
        
        # Combine brown masks
        brown_mask = cv2.bitwise_or(brown_mask1, brown_mask2)
        
        return brown_mask
    
    def _detect_dark_spots_lab(self, lab):
        """Detect dark spots using LAB color space"""
        l_channel = lab[:, :, 0]
        a_channel = lab[:, :, 1]
        
        # Dark areas have low L values
        _, dark_mask = cv2.threshold(l_channel, 80, 255, cv2.THRESH_BINARY_INV)
        
        # Areas with high 'a' values (reddish) combined with darkness
        _, red_mask = cv2.threshold(a_channel, 140, 255, cv2.THRESH_BINARY)
        
        # Combine dark and reddish areas
        lab_mask = cv2.bitwise_and(dark_mask, red_mask)
        
        return lab_mask
    
    def _detect_texture_anomalies(self, image):
        """Detect texture anomalies using local binary patterns concept (calibrated for mango)"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # Calculate local standard deviation (texture measure) - optimized for mango texture
        kernel = np.ones((7, 7), np.float32) / 49  # Slightly smaller kernel for mango details
        gray_float = gray.astype(np.float32)
        mean = cv2.filter2D(gray_float, -1, kernel)
        sqr_mean = cv2.filter2D(gray_float**2, -1, kernel)
        
        # Ensure variance is non-negative (handle floating point precision errors)
        variance = sqr_mean - mean**2
        variance = np.maximum(variance, 0)  # Clamp negative values to 0
        std_dev = np.sqrt(variance)
        
        # Handle NaN and infinity values
        std_dev = np.nan_to_num(std_dev, nan=0.0, posinf=255.0, neginf=0.0)
        
        # Normalize to 0-255 range
        if std_dev.max() > 0:
            std_dev = (std_dev / std_dev.max() * 255).astype(np.uint8)
        else:
            std_dev = std_dev.astype(np.uint8)
        
        # Higher threshold for mango diseases (less sensitive to normal texture)
        _, texture_mask = cv2.threshold(std_dev, 30, 255, cv2.THRESH_BINARY)
        
        # Additional texture analysis for mango-specific roughness
        # Use Sobel operators to detect rough areas - more conservative
        sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3)
        sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3)
        sobel_magnitude = np.sqrt(sobelx**2 + sobely**2)
        sobel_magnitude = np.uint8(sobel_magnitude / sobel_magnitude.max() * 255)
        
        # Higher threshold for rough texture detection (less sensitive)
        _, rough_mask = cv2.threshold(sobel_magnitude, 50, 255, cv2.THRESH_BINARY)
        
        # Combine standard deviation and sobel-based texture detection
        # Use AND operation to require both methods to agree (more conservative)
        final_texture_mask = cv2.bitwise_and(texture_mask, rough_mask)
        
        return final_texture_mask
    
    def _detect_irregular_edges(self, image):
        """Detect irregular edges that might indicate disease boundaries"""
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # Apply Canny edge detection
        edges = cv2.Canny(gray, 50, 150)
        
        # Dilate edges to create regions
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
        dilated_edges = cv2.dilate(edges, kernel, iterations=2)
        
        return dilated_edges
    
    def calculate_disease_severity(self, image, disease_mask, fruit_mask):
        """Calculate disease severity as percentage of affected area"""
        if fruit_mask is not None:
            # Use provided fruit mask
            fruit_area = cv2.countNonZero(fruit_mask)
        else:
            # Fallback: create fruit mask (assuming fruit takes up most of the image)
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            _, fruit_mask = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)
            fruit_area = cv2.countNonZero(fruit_mask)
        
        # Calculate disease area (only within fruit region)
        if fruit_mask is not None:
            disease_in_fruit = cv2.bitwise_and(disease_mask, fruit_mask)
            disease_area = cv2.countNonZero(disease_in_fruit)
        else:
            disease_area = cv2.countNonZero(disease_mask)
        
        if fruit_area == 0:
            return 0
        
        severity_percentage = (disease_area / fruit_area) * 100
        return min(severity_percentage, 100)  # Cap at 100%
    
    def classify_disease_level(self, severity_percentage):
        """Classify mango disease level based on severity percentage (calibrated for mango)"""
        if severity_percentage < 2:
            return "Healthy", (0, 255, 0)  # Green - very strict for healthy
        elif severity_percentage < 8:
            return "Early Disease", (0, 255, 255)  # Yellow - early detection
        elif severity_percentage < 20:
            return "Moderate Disease", (0, 165, 255)  # Orange - moderate infection
        elif severity_percentage < 40:
            return "Severe Disease", (0, 100, 255)  # Red-Orange - severe infection
        else:
            return "Critical Disease", (0, 0, 255)  # Red - critical, unmarketable
    
    def process_image(self, image_path):
        """Main processing function with background removal (calibrated for mango diseases)"""
        # Load image
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError("Could not load image")
        
        print("Step 1: Removing background (mango-optimized)...")
        # Remove background first
        fruit_mask, image_no_bg = self.remove_background(image)
        
        print("Step 2: Detecting mango diseases (Alternaria, Anthracnose, Aspergillus, Lasiodiplodia)...")
        # Detect diseased areas (only within mango region)
        disease_mask = self.detect_diseased_areas(image_no_bg, fruit_mask)
        
        print("Step 3: Calculating mango disease severity...")
        # Calculate severity
        severity = self.calculate_disease_severity(image_no_bg, disease_mask, fruit_mask)
        disease_level, color = self.classify_disease_level(severity)
        
        print("Step 4: Creating mango disease visualization...")
        # Create output visualization with bounding boxes
        output_image, disease_info = self._create_output_visualization(image, disease_mask, severity, disease_level, color, fruit_mask)
        
        # Store results
        self.disease_mask = disease_mask
        self.processed_image = output_image
        self.fruit_mask = fruit_mask
        self.image_no_bg = image_no_bg
        self.disease_info = disease_info
        
        return {
            'severity_percentage': severity,
            'disease_level': disease_level,
            'disease_mask': disease_mask,
            'output_image': output_image,
            'fruit_mask': fruit_mask,
            'image_no_bg': image_no_bg,
            'disease_info': disease_info,
            'num_diseased_regions': len(disease_info)
        }
    
    def _create_output_visualization(self, original, mask, severity, level, color, fruit_mask=None):
        """Create visualization showing detected diseased areas with bounding boxes"""
        # Start with original image
        result = original.copy()
        
        # If we have a fruit mask, dim the background
        if fruit_mask is not None:
            background = np.zeros_like(original)
            result = np.where(fruit_mask[..., np.newaxis] == 0, 
                            background, result)
        
        # Create colored overlay for diseased areas
        overlay = result.copy()
        disease_info = []
        
        if np.any(mask > 0):  # Only if there are diseased areas
            overlay[mask > 0] = color
            
            # Blend original image with overlay
            result = cv2.addWeighted(result, 0.7, overlay, 0.3, 0)
            
            # Find contours and draw bounding boxes
            contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            # Filter contours by minimum area to avoid tiny noise
            min_area = 50  # Minimum area threshold
            valid_contours = [cnt for cnt in contours if cv2.contourArea(cnt) >= min_area]
            
            # Draw contours and bounding boxes
            for i, contour in enumerate(valid_contours):
                # Draw contour outline
                cv2.drawContours(result, [contour], -1, color, 2)
                
                # Get bounding rectangle
                x, y, w, h = cv2.boundingRect(contour)
                
                # Draw bounding box
                cv2.rectangle(result, (x, y), (x + w, y + h), color, 3)
                
                # Calculate area of this diseased region
                area = cv2.contourArea(contour)
                disease_info.append({
                    'id': i + 1,
                    'bbox': (x, y, w, h),
                    'area': area,
                    'center': (x + w//2, y + h//2)
                })
                
                # Add label with disease ID
                label = f"D{i+1}"
                label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                
                # Position label above bounding box
                label_x = x
                label_y = y - 10 if y - 10 > 20 else y + h + 25
                
                # Draw label background
                cv2.rectangle(result, 
                            (label_x - 2, label_y - label_size[1] - 2),
                            (label_x + label_size[0] + 2, label_y + 2),
                            color, -1)
                
                # Draw label text
                cv2.putText(result, label, (label_x, label_y), 
                          cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
        
        # Add overall statistics
        num_diseases = len(disease_info)
        text1 = f"{level}: {severity:.1f}%"
        text2 = f"Diseased Regions: {num_diseases}"
        
        # Main status text
        cv2.putText(result, text1, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, color, 2)
        cv2.putText(result, text2, (10, 70), cv2.FONT_HERSHEY_SIMPLEX, 0.8, color, 2)
        
        # Add fruit outline
        if fruit_mask is not None:
            fruit_contours, _ = cv2.findContours(fruit_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cv2.drawContours(result, fruit_contours, -1, (255, 255, 255), 2)
        
        return result, disease_info

    def save_detailed_report(self, output_path, results):
        """Save detailed report of disease detection results"""
        report_path = output_path.replace('.jpg', '_report.txt').replace('.png', '_report.txt')
        
        with open(report_path, 'w') as f:
            f.write("=== FRUIT DISEASE DETECTION REPORT ===\n")
            f.write(f"Image processed: {output_path}\n")
            f.write(f"Overall Disease Level: {results['disease_level']}\n")
            f.write(f"Overall Severity: {results['severity_percentage']:.2f}%\n")
            f.write(f"Number of Diseased Regions: {results['num_diseased_regions']}\n")
            f.write("\n=== INDIVIDUAL DISEASE REGIONS ===\n")
            
            if results['disease_info']:
                for disease in results['disease_info']:
                    f.write(f"\nDisease Region D{disease['id']}:\n")
                    f.write(f"  - Bounding Box: x={disease['bbox'][0]}, y={disease['bbox'][1]}, ")
                    f.write(f"width={disease['bbox'][2]}, height={disease['bbox'][3]}\n")
                    f.write(f"  - Area: {disease['area']:.0f} pixels\n")
                    f.write(f"  - Center: ({disease['center'][0]}, {disease['center'][1]})\n")
            else:
                f.write("No diseased regions detected.\n")
        
        print(f"Detailed report saved: {report_path}")
        return report_path
    
    def save_results(self, output_path, include_mask=True, include_background_removed=True):
        """Save processing results"""
        if self.processed_image is not None:
            cv2.imwrite(output_path, self.processed_image)
            print(f"Main result saved: {output_path}")
        
        if include_mask and self.disease_mask is not None:
            mask_path = output_path.replace('.', '_disease_mask.')
            cv2.imwrite(mask_path, self.disease_mask)
            print(f"Disease mask saved: {mask_path}")
        
        if include_background_removed and hasattr(self, 'image_no_bg') and self.image_no_bg is not None:
            bg_removed_path = output_path.replace('.', '_no_background.')
            cv2.imwrite(bg_removed_path, self.image_no_bg)
            print(f"Background removed image saved: {bg_removed_path}")
        
        if hasattr(self, 'fruit_mask') and self.fruit_mask is not None:
            fruit_mask_path = output_path.replace('.', '_fruit_mask.')
            cv2.imwrite(fruit_mask_path, self.fruit_mask)
            print(f"Fruit mask saved: {fruit_mask_path}")

# Example usage and testing function
def demonstrate_disease_detection():
    """Demonstrate the mango disease detection algorithm"""
    detector = FruitDiseaseDetector()
    
    print("Mango Disease Detection Algorithm with Background Removal & Bounding Boxes")
    print("========================================================================")
    print("This algorithm detects mango diseases using:")
    print("1. Mango-optimized background removal (GrabCut + Color + Edge detection)")
    print("2. Anthracnose detection (dark circular spots with orange/pink halos)")
    print("3. Alternaria detection (brown/black irregular spots with concentric rings)")
    print("4. Aspergillus detection (black mould with greenish tints)")
    print("5. Lasiodiplodia detection (stem rot with brown/black soft areas)")
    print("6. Mango-calibrated texture analysis for rough surfaces")
    print("7. Edge detection for irregular disease boundaries")
    print("8. Individual disease region bounding boxes")
    print()
    
    print("Algorithm Features for Mango:")
    print("- Multi-method background removal optimized for mango shape")
    print("- Disease-specific detection for common mango diseases")
    print("- Individual region labeling (D1, D2, D3, etc.)")
    print("- Mango-specific severity assessment (Healthy < 2%, Early < 8%, etc.)")
    print("- Comprehensive reporting for mango disease management")
    print("- Calibrated for mango color variations (green, yellow, orange, red)")
    print("- Specialized detection for Aspergillus (Black Mould) and Lasiodiplodia (Stem Rot)")
    
    return detector

# Batch testing function for algorithm validation
def test_mango_detection_algorithm():
    """Test the algorithm on multiple mango samples"""
    detector = FruitDiseaseDetector()
    
    # Test cases with expected results - using correct file names
    test_cases = [
        ("SenMangoFruitDDS_bgremoved/Healthy/healthy_003.jpg", "Should be Healthy"),
        ("SenMangoFruitDDS_bgremoved/Healthy/healthy_010.jpg", "Should be Healthy"),
        ("SenMangoFruitDDS_bgremoved/Alternaria/Alternaria_005.jpg", "Should detect Alternaria"),
        ("SenMangoFruitDDS_bgremoved/Alternaria/Alternaria_010.jpg", "Should detect Alternaria"),
        ("SenMangoFruitDDS_bgremoved/Anthracnose/Anthracnose_002.jpg", "Should detect Anthracnose"),
        ("SenMangoFruitDDS_bgremoved/Anthracnose/Anthracnose_010.jpg", "Should detect Anthracnose"),
        ("SenMangoFruitDDS_bgremoved/Black Mould Rot/Aspergillus_001.jpg", "Should detect Aspergillus (Black Mould)"),
        ("SenMangoFruitDDS_bgremoved/Black Mould Rot/Aspergillus_010.jpg", "Should detect Aspergillus (Black Mould)"),
        ("SenMangoFruitDDS_bgremoved/Stem and Rot/Lasiodiplodia_001.jpg", "Should detect Lasiodiplodia (Stem Rot)"),
        ("SenMangoFruitDDS_bgremoved/Stem and Rot/Lasiodiplodia_012.jpg", "Should detect Lasiodiplodia (Stem Rot)"),
    ]
    
    print("=== MANGO DISEASE DETECTION ALGORITHM VALIDATION ===")
    print("Testing multiple samples to validate algorithm performance")
    print("Diseases: Alternaria, Anthracnose, Aspergillus (Black Mould), Lasiodiplodia (Stem Rot)")
    print("=" * 80)
    
    results_summary = []
    
    for i, (image_path, expected) in enumerate(test_cases, 1):
        print(f"\nTest {i}: {image_path}")
        print(f"Expected: {expected}")
        
        try:
            results = detector.process_image(image_path)
            output_path = f"test_result_{i}.jpg"
            
            result_info = {
                'test_id': i,
                'image_path': image_path,
                'expected': expected,
                'detected_level': results['disease_level'],
                'severity': results['severity_percentage'],
                'num_regions': results['num_diseased_regions']
            }
            results_summary.append(result_info)
            
            print(f"Result: {results['disease_level']} ({results['severity_percentage']:.2f}%)")
            print(f"Regions: {results['num_diseased_regions']}")
            
            # Save test result
            detector.save_results(output_path, include_mask=True)
            
        except Exception as e:
            print(f"Error: {e}")
            results_summary.append({
                'test_id': i,
                'image_path': image_path,
                'expected': expected,
                'detected_level': 'ERROR',
                'severity': 0.0,
                'num_regions': 0
            })
        
        print("-" * 50)
    
    # Print summary
    print("\n" + "=" * 80)
    print("TESTING SUMMARY")
    print("=" * 80)
    for result in results_summary:
        status = "✓" if result['detected_level'] != 'Healthy' and 'Should detect' in result['expected'] else "✗" if result['detected_level'] == 'Healthy' and 'Should detect' in result['expected'] else "✓"
        print(f"Test {result['test_id']:2d}: {status} {result['detected_level']} ({result['severity']:.1f}%) - {result['expected']}")
    
    return results_summary

# Usage example:
if __name__ == "__main__":
    # Initialize mango disease detector
    detector = FruitDiseaseDetector()
    
    # Test different disease types - uncomment to test specific diseases:
    # image_path = "SenMangoFruitDDS_bgremoved/Healthy/healthy_003.jpg"  # Should be Healthy
    # image_path = "SenMangoFruitDDS_bgremoved/Alternaria/Alternaria_005.jpg"  # Should detect Alternaria
    # image_path = "SenMangoFruitDDS_bgremoved/Anthracnose/Anthracnose_002.jpg"  # Should detect Anthracnose
    # image_path = "SenMangoFruitDDS_bgremoved/Black Mould Rot/Aspergillus_001.jpg"  # Should detect Aspergillus
    image_path = "SenMangoFruitDDS_bgremoved/Stem and Rot/Lasiodiplodia_012.jpg"  # Should detect Lasiodiplodia
    
    output_path = "mango_disease_detection_result.jpg"  # Output file name
    
    try:
        print("Processing mango image with calibrated algorithm...")
        print(f"Input: {image_path}")
        print(f"Output: {output_path}")
        print("-" * 60)
        
        results = detector.process_image(image_path)
        
        print(f"\n=== MANGO DISEASE DETECTION RESULTS ===")
        print(f"Disease Level: {results['disease_level']}")
        print(f"Severity: {results['severity_percentage']:.2f}%")
        print(f"Number of Diseased Regions: {results['num_diseased_regions']}")
        
        # Print individual disease information
        if results['disease_info']:
            print(f"\n=== INDIVIDUAL DISEASE REGIONS ===")
            for disease in results['disease_info']:
                print(f"Disease D{disease['id']}: Area={disease['area']:.0f} pixels, "
                      f"Center=({disease['center'][0]}, {disease['center'][1]})")
        else:
            print("\n=== MANGO HEALTH STATUS ===")
            print("No significant disease regions detected - mango appears healthy!")
        
        # Save results
        detector.save_results(output_path, include_mask=True)
        
        # Save detailed report
        detector.save_detailed_report(output_path, results)
        
        print(f"\n=== FILES SAVED ===")
        print(f"Main result: {output_path}")
        print(f"Disease mask: {output_path.replace('.', '_disease_mask.')}")
        print(f"Background removed: {output_path.replace('.', '_no_background.')}")
        print(f"Fruit mask: {output_path.replace('.', '_fruit_mask.')}")
        print(f"Detailed report: {output_path.replace('.jpg', '_report.txt')}")
        
        print("\nMango disease detection completed successfully!")
        print("The algorithm has been calibrated for mango-specific diseases:")
        print("- Alternaria, Anthracnose, Aspergillus (Black Mould), Lasiodiplodia (Stem Rot)")
        
    except FileNotFoundError:
        print(f"Error: Could not find mango image file '{image_path}'")
        print("Please check the file path and make sure the image exists.")
        print("\nAvailable test images:")
        print("- SenMangoFruitDDS_bgremoved/Healthy/healthy_003.jpg")
        print("- SenMangoFruitDDS_bgremoved/Alternaria/Alternaria_005.jpg")
        print("- SenMangoFruitDDS_bgremoved/Anthracnose/Anthracnose_002.jpg")
    except Exception as e:
        print(f"Error processing mango image: {e}")
    
    # Uncomment below to see algorithm demonstration
    # demonstrate_disease_detection()
    
    # Uncomment below to run batch testing
    # test_mango_detection_algorithm()