Update part3.py

part3.py

@@ -7,14 +7,12 @@ from PIL import Image
 
 class SAMAnalyzer:
     def __init__(self, model_path="sam_vit_h_4b8939.pth"):
+        """Initialize SAM model for farmland segmentation"""
         self.model_path = model_path
         self.sam = None
         self.predictor = None
-        self.initialize_sam()
-
-    def initialize_sam(self):
-        """Initialize SAM model"""
         try:
+            print("Initializing SAM model...")
             self.sam = sam_model_registry["vit_h"](checkpoint=self.model_path)
             self.predictor = SamPredictor(self.sam)
             print("SAM model initialized successfully")
@@ -22,158 +20,191 @@ class SAMAnalyzer:
             print(f"Error initializing SAM model: {e}")
             raise
 
-    def process_image(self, image_data):
+    def process_image(self, image):
         """Process uploaded image using SAM"""
         try:
-            # Convert uploaded image to numpy array
-            if isinstance(image_data, (str, bytes)):
-                if isinstance(image_data, str):
-                    image = cv2.imread(image_data)
-                else:
-                    nparr = np.frombuffer(image_data, np.uint8)
-                    image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            print("Starting image processing...")
+            if image is None:
+                raise ValueError("No image provided")
+
+            # Ensure image is in correct format
+            if isinstance(image, np.ndarray):
+                if len(image.shape) == 2:  # Grayscale
+                    image = np.stack((image,)*3, axis=-1)
+                elif len(image.shape) == 3 and image.shape[2] == 4:  # RGBA
+                    image = image[:,:,:3]
             else:
-                image = np.array(Image.open(image_data))
+                raise ValueError("Invalid image format")
 
-            # Segment farmland
             print("Segmenting farmland...")
             farmland_mask = self.segment_farmland(image)
 
-            # Calculate vegetation index
             print("Calculating vegetation index...")
             veg_index = self.calculate_vegetation_index(image, farmland_mask)
 
-            # Analyze health
             print("Analyzing crop health...")
             health_analysis = self.analyze_crop_health(veg_index, farmland_mask)
 
-            # Create visualization
-            print("Generating visualization...")
+            print("Creating visualization...")
             viz_plot = self.create_visualization(image, farmland_mask, veg_index)
 
             return veg_index, health_analysis, viz_plot
 
         except Exception as e:
-            print(f"Error processing image: {e}")
+            print(f"Error in image processing: {e}")
             return None, None, None
 
     def segment_farmland(self, image):
         """Segment farmland using SAM2"""
-        self.predictor.set_image(image)
-
-        # Generate automatic mask proposals
-        h, w = image.shape[:2]
-        input_point = np.array([[w//2, h//2]])
-        input_label = np.array([1])
-
-        masks, scores, logits = self.predictor.predict(
-            point_coords=input_point,
-            point_labels=input_label,
-            multimask_output=True
-        )
+        try:
+            self.predictor.set_image(image)
+            
+            # Generate automatic mask proposals
+            h, w = image.shape[:2]
+            input_point = np.array([[w//2, h//2]])  # Center point
+            input_label = np.array([1])  # Foreground
+
+            masks, scores, logits = self.predictor.predict(
+                point_coords=input_point,
+                point_labels=input_label,
+                multimask_output=True
+            )
+
+            # Select best mask
+            best_mask = masks[scores.argmax()]
+            return best_mask
 
-        # Select best mask
-        best_mask = masks[scores.argmax()]
-        return best_mask
+        except Exception as e:
+            print(f"Error in farmland segmentation: {e}")
+            raise
 
     def calculate_vegetation_index(self, image, mask):
         """Calculate vegetation index using RGB"""
-        r, g, b = image[:,:,0], image[:,:,1], image[:,:,2]
+        try:
+            # Extract RGB channels
+            r, g, b = image[:,:,0], image[:,:,1], image[:,:,2]
+
+            # Calculate visible-band vegetation index
+            numerator = (2 * g.astype(float) - r.astype(float) - b.astype(float))
+            denominator = (2 * g.astype(float) + r.astype(float) + b.astype(float))
+            
+            # Avoid division by zero
+            denominator[denominator == 0] = 1e-10
 
-        numerator = (2 * g.astype(float) - r.astype(float) - b.astype(float))
-        denominator = (2 * g.astype(float) + r.astype(float) + b.astype(float))
-        denominator[denominator == 0] = 1e-10
+            # Calculate index and normalize
+            veg_index = numerator / denominator
+            veg_index = (veg_index + 1) / 2  # Normalize to 0-1 range
+            
+            # Apply mask
+            veg_index = veg_index * mask
 
-        veg_index = numerator / denominator
-        veg_index = (veg_index + 1) / 2
-        veg_index = veg_index * mask
+            return veg_index
 
-        return veg_index
+        except Exception as e:
+            print(f"Error calculating vegetation index: {e}")
+            raise
 
     def analyze_crop_health(self, veg_index, mask):
         """Analyze crop health based on vegetation index"""
-        valid_pixels = veg_index[mask > 0]
-        if len(valid_pixels) == 0:
-            return {
-                'average_index': 0,
-                'health_distribution': {
-                    'low_vegetation': 0,
-                    'moderate_vegetation': 0,
-                    'high_vegetation': 0
-                },
-                'overall_health': 'No vegetation detected'
+        try:
+            valid_pixels = veg_index[mask > 0]
+            if len(valid_pixels) == 0:
+                return {
+                    'average_index': 0,
+                    'health_distribution': {
+                        'low_vegetation': 0,
+                        'moderate_vegetation': 0,
+                        'high_vegetation': 0
+                    },
+                    'overall_health': 'No vegetation detected'
+                }
+
+            avg_index = np.mean(valid_pixels)
+            health_categories = {
+                'low_vegetation': np.sum((valid_pixels <= 0.3)) / len(valid_pixels),
+                'moderate_vegetation': np.sum((valid_pixels > 0.3) & (valid_pixels <= 0.6)) / len(valid_pixels),
+                'high_vegetation': np.sum((valid_pixels > 0.6)) / len(valid_pixels)
             }
 
-        avg_index = np.mean(valid_pixels)
-        health_categories = {
-            'low_vegetation': np.sum((valid_pixels <= 0.3)) / len(valid_pixels),
-            'moderate_vegetation': np.sum((valid_pixels > 0.3) & (valid_pixels <= 0.6)) / len(valid_pixels),
-            'high_vegetation': np.sum((valid_pixels > 0.6)) / len(valid_pixels)
-        }
+            overall_health = 'Healthy' if avg_index > 0.5 else 'Needs attention'
 
-        return {
-            'average_index': avg_index,
-            'health_distribution': health_categories,
-            'overall_health': 'Healthy' if avg_index > 0.5 else 'Needs attention'
-        }
+            return {
+                'average_index': avg_index,
+                'health_distribution': health_categories,
+                'overall_health': overall_health
+            }
+
+        except Exception as e:
+            print(f"Error analyzing crop health: {e}")
+            raise
 
     def create_visualization(self, image, mask, veg_index):
         """Create visualization of results"""
-        fig = plt.figure(figsize=(15, 5))
-
-        # Original image with mask overlay
-        plt.subplot(131)
-        plt.imshow(image)
-        plt.imshow(mask, alpha=0.3, cmap='gray')
-        plt.title('Segmented Farmland')
-        plt.axis('off')
-
-        # Vegetation index heatmap
-        plt.subplot(132)
-        plt.imshow(veg_index, cmap='RdYlGn')
-        plt.colorbar(label='Vegetation Index')
-        plt.title('Vegetation Index')
-        plt.axis('off')
-
-        # Health classification
-        plt.subplot(133)
-        health_mask = np.zeros_like(veg_index)
-        health_mask[veg_index <= 0.3] = 1
-        health_mask[(veg_index > 0.3) & (veg_index <= 0.6)] = 2
-        health_mask[veg_index > 0.6] = 3
-        health_mask = health_mask * mask
-        plt.imshow(health_mask, cmap='viridis')
-        plt.colorbar(ticks=[1, 2, 3], 
-                    label='Vegetation Levels',
-                    boundaries=np.arange(0.5, 4.5),
-                    values=[1, 2, 3])
-        plt.title('Vegetation Levels')
-        plt.axis('off')
-
-        plt.tight_layout()
-        
-        # Save plot to buffer
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png', bbox_inches='tight')
-        buf.seek(0)
-        plt.close()
-        
-        return buf
+        try:
+            fig = plt.figure(figsize=(15, 5))
+
+            # Original image with mask overlay
+            plt.subplot(131)
+            plt.imshow(image)
+            plt.imshow(mask, alpha=0.3, cmap='gray')
+            plt.title('Segmented Farmland')
+            plt.axis('off')
+
+            # Vegetation index heatmap
+            plt.subplot(132)
+            plt.imshow(veg_index, cmap='RdYlGn')
+            plt.colorbar(label='Vegetation Index')
+            plt.title('Vegetation Index')
+            plt.axis('off')
+
+            # Health classification
+            plt.subplot(133)
+            health_mask = np.zeros_like(veg_index)
+            health_mask[veg_index <= 0.3] = 1  # Low
+            health_mask[(veg_index > 0.3) & (veg_index <= 0.6)] = 2  # Moderate
+            health_mask[veg_index > 0.6] = 3  # High
+            health_mask = health_mask * mask
+            plt.imshow(health_mask, cmap='viridis')
+            plt.colorbar(
+                ticks=[1, 2, 3],
+                label='Vegetation Levels',
+                boundaries=np.arange(0.5, 4.5),
+                values=[1, 2, 3]
+            )
+            plt.title('Vegetation Levels')
+            plt.axis('off')
+
+            plt.tight_layout()
+            
+            # Save plot to buffer
+            buf = io.BytesIO()
+            plt.savefig(buf, format='png', bbox_inches='tight', dpi=300)
+            buf.seek(0)
+            plt.close()
+            
+            return buf
+
+        except Exception as e:
+            print(f"Error creating visualization: {e}")
+            raise
 
     def format_analysis_text(self, health_analysis):
         """Format health analysis results as text"""
-        return f"""
-        🌿 Vegetation Analysis Results:
-        
-        📊 Average Vegetation Index: {health_analysis['average_index']:.2f}
-        
-        🌱 Vegetation Distribution:
-        • Low Vegetation: {health_analysis['health_distribution']['low_vegetation']*100:.1f}%
-        • Moderate Vegetation: {health_analysis['health_distribution']['moderate_vegetation']*100:.1f}%
-        • High Vegetation: {health_analysis['health_distribution']['high_vegetation']*100:.1f}%
-        
-        📋 Overall Health Status: {health_analysis['overall_health']}
-        
-        Note: Analysis uses SAM2 for farmland segmentation
-        """
+        try:
+            return f"""
+            🌿 Vegetation Analysis Results:
+            
+            📊 Average Vegetation Index: {health_analysis['average_index']:.2f}
+            
+            🌱 Vegetation Distribution:
+            • Low Vegetation: {health_analysis['health_distribution']['low_vegetation']*100:.1f}%
+            • Moderate Vegetation: {health_analysis['health_distribution']['moderate_vegetation']*100:.1f}%
+            • High Vegetation: {health_analysis['health_distribution']['high_vegetation']*100:.1f}%
+            
+            📋 Overall Health Status: {health_analysis['overall_health']}
+            
+            Note: Analysis uses SAM2 for farmland segmentation
+            """
+        except Exception as e:
+            print(f"Error formatting analysis text: {e}")
+            return "Error generating analysis report"