Update part3.py

part3.py

@@ -2,7 +2,6 @@ import numpy as np
 import cv2
 from segment_anything import sam_model_registry, SamPredictor
 import matplotlib.pyplot as plt
-import io
 from PIL import Image
 
 class SAMAnalyzer:
@@ -33,6 +32,9 @@ class SAMAnalyzer:
                     image = np.stack((image,)*3, axis=-1)
                 elif len(image.shape) == 3 and image.shape[2] == 4:  # RGBA
                     image = image[:,:,:3]
+                # Ensure image is in RGB format
+                if image.shape[2] == 3:
+                    image = cv2.cvtColor(cv2.cvtColor(image, cv2.COLOR_RGB2BGR), cv2.COLOR_BGR2RGB)
             else:
                 raise ValueError("Invalid image format")
 
@@ -46,9 +48,9 @@ class SAMAnalyzer:
             health_analysis = self.analyze_crop_health(veg_index, farmland_mask)
 
             print("Creating visualization...")
-            viz_plot = self.create_visualization(image, farmland_mask, veg_index)
+            fig = self.create_visualization(image, farmland_mask, veg_index)
 
-            return veg_index, health_analysis, viz_plot
+            return veg_index, health_analysis, fig
 
         except Exception as e:
             print(f"Error in image processing: {e}")
@@ -71,8 +73,11 @@ class SAMAnalyzer:
             )
 
             # Select best mask
-            best_mask = masks[scores.argmax()]
-            return best_mask
+            if len(masks) > 0:
+                best_mask = masks[scores.argmax()]
+                return best_mask
+            else:
+                raise ValueError("No valid masks generated")
 
         except Exception as e:
             print(f"Error in farmland segmentation: {e}")
@@ -141,6 +146,10 @@ class SAMAnalyzer:
     def create_visualization(self, image, mask, veg_index):
         """Create visualization of results"""
         try:
+            # Clear any existing plots
+            plt.close('all')
+            
+            # Create figure
             fig = plt.figure(figsize=(15, 5))
 
             # Original image with mask overlay
@@ -152,8 +161,8 @@ class SAMAnalyzer:
 
             # Vegetation index heatmap
             plt.subplot(132)
-            plt.imshow(veg_index, cmap='RdYlGn')
-            plt.colorbar(label='Vegetation Index')
+            im = plt.imshow(veg_index, cmap='RdYlGn', vmin=0, vmax=1)
+            plt.colorbar(im, label='Vegetation Index')
             plt.title('Vegetation Index')
             plt.axis('off')
 
@@ -164,47 +173,19 @@ class SAMAnalyzer:
             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]
-            )
+            
+            im = plt.imshow(health_mask, cmap='viridis', vmin=1, vmax=3)
+            cbar = plt.colorbar(im, ticks=[1, 2, 3])
+            cbar.set_label('Vegetation Levels')
+            cbar.set_ticklabels(['Low', 'Moderate', 'High'])
             plt.title('Vegetation Levels')
             plt.axis('off')
 
+            # Adjust layout
             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
+            return fig
 
         except Exception as e:
             print(f"Error creating visualization: {e}")
-            raise
-
-    def format_analysis_text(self, health_analysis):
-        """Format health analysis results as text"""
-        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"
+            raise