Update app.py

fixed app module

app.py

@@ -5,7 +5,16 @@ Provides a web interface for making predictions with StableResNet
 Author: najahpokkiri
 Date: 2025-05-17
 """
+"""
+Biomass Prediction Gradio App
+Author: najahpokkiri
+Date: 2025-05-17
+
+This app allows users to predict above-ground biomass from satellite imagery
+using a trained StableResNet model.
+"""
 import os
+import sys
 import torch
 import numpy as np
 import gradio as gr
@@ -15,33 +24,48 @@ import matplotlib.pyplot as plt
 import matplotlib.colors as colors
 from PIL import Image
 import io
+import logging
 from huggingface_hub import hf_hub_download
 
+# Configure logger
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
 # Import model architecture
 from model import StableResNet
 
 class BiomassPredictorApp:
-    """Gradio app for biomass prediction"""
+    """Gradio app for biomass prediction from satellite imagery"""
     
     def __init__(self, model_repo="pokkiri/biomass-model"):
+        """Initialize the app with model repository information"""
         self.model = None
         self.package = None
+        self.feature_names = []
         self.model_repo = model_repo
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         
+        # Cache for storing temporary files
+        self.temp_files = []
+        
         # Load the model
         self.load_model()
     
     def load_model(self):
         """Load the model and preprocessing pipeline from HuggingFace Hub"""
         try:
-            # Download files from HuggingFace
+            logger.info(f"Loading model from {self.model_repo}")
+            
+            # Download model files from HuggingFace
             model_path = hf_hub_download(repo_id=self.model_repo, filename="model.pt")
             package_path = hf_hub_download(repo_id=self.model_repo, filename="model_package.pkl")
             
             # Load package with metadata
             self.package = joblib.load(package_path)
+            
+            # Extract information from package
             n_features = self.package['n_features']
+            self.feature_names = self.package.get('feature_names', [f"feature_{i}" for i in range(n_features)])
             
             # Initialize model
             self.model = StableResNet(n_features=n_features)
@@ -49,17 +73,33 @@ class BiomassPredictorApp:
             self.model.to(self.device)
             self.model.eval()
             
-            print(f"Model loaded successfully from {self.model_repo}")
-            print(f"Number of features: {n_features}")
-            print(f"Using device: {self.device}")
+            logger.info(f"Model loaded successfully from {self.model_repo}")
+            logger.info(f"Number of features: {n_features}")
+            logger.info(f"Using device: {self.device}")
             
             return True
         except Exception as e:
-            print(f"Error loading model: {e}")
+            logger.error(f"Error loading model: {e}")
+            import traceback
+            logger.error(traceback.format_exc())
             return False
     
+    def cleanup(self):
+        """Clean up temporary files"""
+        for tmp_path in self.temp_files:
+            try:
+                if os.path.exists(tmp_path):
+                    os.unlink(tmp_path)
+            except Exception as e:
+                logger.warning(f"Failed to remove temporary file {tmp_path}: {e}")
+        
+        self.temp_files = []
+    
     def predict_biomass(self, image_file, display_type="heatmap"):
         """Predict biomass from a satellite image"""
+        if self.model is None:
+            return None, "Error: Model not loaded. Please check logs for details."
+        
         try:
             # Create a temporary file to save the uploaded file
             with tempfile.NamedTemporaryFile(suffix='.tif', delete=False) as tmp_file:
@@ -67,10 +107,14 @@ class BiomassPredictorApp:
                 with open(image_file.name, 'rb') as f:
                     tmp_file.write(f.read())
             
+            # Add to list for cleanup later
+            self.temp_files.append(tmp_path)
+            
+            # Ensure rasterio is available
             try:
                 import rasterio
             except ImportError:
-                return None, "Error: rasterio is required but not installed."
+                return None, "Error: rasterio is required but not installed. Please install with: pip install rasterio"
             
             # Open the image file
             with rasterio.open(tmp_path) as src:
@@ -79,29 +123,39 @@ class BiomassPredictorApp:
                 transform = src.transform
                 crs = src.crs
                 
-                # Check if number of bands matches expected features
+                # Validate image dimensions
                 if image.shape[0] < self.package['n_features']:
-                    return None, f"Error: Image has {image.shape[0]} bands, but model expects at least {self.package['n_features']} features."
+                    return None, (f"Error: Image has {image.shape[0]} bands, but model expects at least "
+                                 f"{self.package['n_features']} features.")
                 
-                print(f"Processing image: {height}x{width} pixels, {image.shape[0]} bands")
+                logger.info(f"Processing image: {height}x{width} pixels, {image.shape[0]} bands")
                 
                 # Process in chunks to avoid memory issues
-                chunk_size = 1000
+                chunk_size = min(1000, height, width)  # Adjust chunk size for smaller images
                 predictions = np.zeros((height, width), dtype=np.float32)
                 
                 # Create mask for valid pixels (not NaN or Inf)
                 valid_mask = np.all(np.isfinite(image), axis=0)
                 
+                # Show progress indicator
+                progress_text = f"Processing {height}x{width} image..."
+                logger.info(progress_text)
+                
                 # Process image in chunks
+                total_chunks = ((height + chunk_size - 1) // chunk_size) * ((width + chunk_size - 1) // chunk_size)
+                chunk_count = 0
+                
                 for y_start in range(0, height, chunk_size):
                     y_end = min(y_start + chunk_size, height)
                     
                     for x_start in range(0, width, chunk_size):
                         x_end = min(x_start + chunk_size, width)
+                        chunk_count += 1
                         
                         # Get chunk mask
                         chunk_mask = valid_mask[y_start:y_end, x_start:x_end]
                         if not np.any(chunk_mask):
+                            logger.info(f"Skipping chunk {chunk_count}/{total_chunks} (no valid pixels)")
                             continue
                         
                         # Extract valid pixels
@@ -124,42 +178,58 @@ class BiomassPredictorApp:
                             batch_predictions = self.model(batch_tensor).cpu().numpy()
                         
                         # Convert from log scale if needed
-                        if self.package['use_log_transform']:
-                            batch_predictions = np.exp(batch_predictions) - self.package.get('epsilon', 1.0)
+                        if self.package.get('use_log_transform', False):
+                            epsilon = self.package.get('epsilon', 1.0)
+                            batch_predictions = np.exp(batch_predictions) - epsilon
                             batch_predictions = np.maximum(batch_predictions, 0)  # Ensure non-negative
                         
                         # Insert predictions back into the image
                         for idx, (i, j) in enumerate(zip(valid_y, valid_x)):
                             predictions[y_start+i, x_start+j] = batch_predictions[idx]
-                
-                # Delete temporary file
-                os.unlink(tmp_path)
+                        
+                        logger.info(f"Processed chunk {chunk_count}/{total_chunks}")
                 
                 # Create visualization
+                logger.info("Creating visualization...")
                 plt.figure(figsize=(12, 8))
                 
                 if display_type == "heatmap":
                     # Create heatmap
-                    plt.imshow(predictions, cmap='viridis')
+                    # Use masked array for better visualization
+                    masked_predictions = np.ma.masked_where(~valid_mask, predictions)
+                    
+                    # Set min/max values based on percentiles for better contrast
+                    vmin = np.percentile(predictions[valid_mask], 1)
+                    vmax = np.percentile(predictions[valid_mask], 99)
+                    
+                    plt.imshow(masked_predictions, cmap='viridis', vmin=vmin, vmax=vmax)
                     plt.colorbar(label='Biomass (Mg/ha)')
                     plt.title('Predicted Above-Ground Biomass')
+                    plt.axis('off')  # Hide axes for cleaner visualization
                     
                 elif display_type == "rgb_overlay":
                     # Create RGB + overlay
                     if image.shape[0] >= 3:
                         # Use first 3 bands as RGB
                         rgb = image[[0, 1, 2]].transpose(1, 2, 0)
-                        rgb = np.clip((rgb - np.percentile(rgb, 2)) / (np.percentile(rgb, 98) - np.percentile(rgb, 2)), 0, 1)
                         
-                        plt.imshow(rgb)
+                        # Enhance contrast with percentile-based normalization
+                        p2 = np.percentile(rgb[np.isfinite(rgb)], 2)
+                        p98 = np.percentile(rgb[np.isfinite(rgb)], 98)
+                        rgb_norm = np.clip((rgb - p2) / (p98 - p2), 0, 1)
+                        
+                        # Display RGB image
+                        plt.imshow(rgb_norm)
                         
                         # Create mask for overlay (where we have predictions)
-                        mask = ~np.isclose(predictions, 0)
+                        mask = valid_mask & (~np.isclose(predictions, 0))
                         overlay = np.zeros((height, width, 4))
                         
                         # Create colormap for biomass
-                        norm = colors.Normalize(vmin=np.percentile(predictions[mask], 5), 
-                                              vmax=np.percentile(predictions[mask], 95))
+                        norm = colors.Normalize(
+                            vmin=np.percentile(predictions[mask], 5), 
+                            vmax=np.percentile(predictions[mask], 95)
+                        )
                         cmap = plt.cm.viridis
                         
                         # Apply colormap
@@ -170,42 +240,63 @@ class BiomassPredictorApp:
                         plt.colorbar(plt.cm.ScalarMappable(norm=norm, cmap=cmap), 
                                     label='Biomass (Mg/ha)')
                         plt.title('Biomass Prediction Overlay')
+                        plt.axis('off')
                     else:
-                        plt.imshow(predictions, cmap='viridis')
+                        # Fallback to regular heatmap if not enough bands for RGB
+                        masked_predictions = np.ma.masked_where(~valid_mask, predictions)
+                        plt.imshow(masked_predictions, cmap='viridis')
                         plt.colorbar(label='Biomass (Mg/ha)')
                         plt.title('Predicted Above-Ground Biomass')
+                        plt.axis('off')
                 
                 # Save figure to bytes buffer
                 buf = io.BytesIO()
                 plt.savefig(buf, format='png', dpi=150, bbox_inches='tight')
                 buf.seek(0)
+                plt.close()
                 
-                # Create summary statistics
+                # Calculate summary statistics
                 valid_predictions = predictions[valid_mask]
                 stats = {
                     'Mean Biomass': f"{np.mean(valid_predictions):.2f} Mg/ha",
                     'Median Biomass': f"{np.median(valid_predictions):.2f} Mg/ha",
                     'Min Biomass': f"{np.min(valid_predictions):.2f} Mg/ha",
-                    'Max Biomass': f"{np.max(valid_predictions):.2f} Mg/ha",
-                    'Total Biomass': f"{np.sum(valid_predictions) * (transform[0] * transform[0]) / 10000:.2f} Mg",
-                    'Area': f"{np.sum(valid_mask) * (transform[0] * transform[0]) / 10000:.2f} hectares"
+                    'Max Biomass': f"{np.max(valid_predictions):.2f} Mg/ha"
                 }
                 
+                # Add area and total biomass if transform is available
+                if transform is not None:
+                    pixel_area_m2 = abs(transform[0] * transform[4])  # Assuming square pixels
+                    total_biomass = np.sum(valid_predictions) * (pixel_area_m2 / 10000)  # Convert to hectares
+                    area_hectares = np.sum(valid_mask) * (pixel_area_m2 / 10000)
+                    
+                    stats['Total Biomass'] = f"{total_biomass:.2f} Mg"
+                    stats['Area'] = f"{area_hectares:.2f} hectares"
+                
                 # Format statistics as markdown
                 stats_md = "### Biomass Statistics\n\n"
                 stats_md += "| Metric | Value |\n|--------|-------|\n"
                 for k, v in stats.items():
                     stats_md += f"| {k} | {v} |\n"
                 
-                # Close the plot
-                plt.close()
+                # Add processing info
+                stats_md += f"\n\n*Processed {np.sum(valid_mask):,} valid pixels*"
+                
+                # Cleanup temporary files
+                self.cleanup()
                 
                 # Return visualization and statistics
                 return Image.open(buf), stats_md
                 
         except Exception as e:
+            # Ensure cleanup even on error
+            self.cleanup()
+            
             import traceback
-            return None, f"Error predicting biomass: {str(e)}\n\n{traceback.format_exc()}"
+            logger.error(f"Error predicting biomass: {e}")
+            logger.error(traceback.format_exc())
+            
+            return None, f"Error predicting biomass: {str(e)}\n\nPlease check logs for details."
 
     def create_interface(self):
         """Create Gradio interface"""
@@ -220,7 +311,7 @@ class BiomassPredictorApp:
             """)
             
             with gr.Row():
-                with gr.Column():
+                with gr.Column(scale=1):
                     input_image = gr.File(
                         label="Upload Satellite Image (GeoTIFF)",
                         file_types=[".tif", ".tiff"]
@@ -232,9 +323,9 @@ class BiomassPredictorApp:
                         label="Display Type"
                     )
                     
-                    submit_btn = gr.Button("Generate Biomass Prediction")
+                    submit_btn = gr.Button("Generate Biomass Prediction", variant="primary")
                 
-                with gr.Column():
+                with gr.Column(scale=2):
                     output_image = gr.Image(
                         label="Biomass Prediction Map",
                         type="pil"
@@ -245,7 +336,7 @@ class BiomassPredictorApp:
                     )
             
             with gr.Accordion("About", open=False):
-                gr.Markdown(f"""
+                gr.Markdown("""
                 ## About This Model
                 
                 This biomass prediction model uses the StableResNet architecture to predict above-ground biomass from satellite imagery.
@@ -255,9 +346,9 @@ class BiomassPredictorApp:
                 - Architecture: StableResNet
                 - Input: Multi-spectral satellite imagery
                 - Output: Above-ground biomass (Mg/ha)
-                - Creator: {pokkiri}
-                - Date: {2025-05-17}
-                - Model Repository: [{pokkiri/biomass-model}](https://huggingface.co/{pokkiri/biomass-model})
+                - Creator: najahpokkiri
+                - Date: 2025-05-17
+                - Model Repository: [pokkiri/biomass-model](https://huggingface.co/pokkiri/biomass-model)
                 
                 ### How It Works
                 
@@ -267,6 +358,27 @@ class BiomassPredictorApp:
                 4. Results are visualized as a heatmap or RGB overlay
                 """)
             
+            with gr.Accordion("Examples", open=False):
+                gr.Markdown("""
+                ### Example Data
+                
+                To try the model, you can use sample GeoTIFF files with the following characteristics:
+                
+                - Multi-band satellite imagery (Sentinel-2, Landsat, etc.)
+                - Contains bands in the proper order (see documentation)
+                - Images should be relatively small (< 2000x2000 pixels) for faster processing
+                
+                You can find sample data at:
+                - [Earth Explorer](https://earthexplorer.usgs.gov/)
+                - [Copernicus Open Access Hub](https://scihub.copernicus.eu/)
+                - [Planetary Computer](https://planetarycomputer.microsoft.com/)
+                """)
+            
+            # Add a warning if model failed to load
+            if self.model is None:
+                gr.Warning("⚠️ Model failed to load. The app may not work correctly. Check logs for details.")
+            
+            # Connect the submit button
             submit_btn.click(
                 fn=self.predict_biomass,
                 inputs=[input_image, display_type],
@@ -277,9 +389,19 @@ class BiomassPredictorApp:
 
 def launch_app():
     """Launch the Gradio app"""
-    app = BiomassPredictorApp()
-    interface = app.create_interface()
-    interface.launch()
+    try:
+        # Create app instance
+        app = BiomassPredictorApp()
+        
+        # Create interface
+        interface = app.create_interface()
+        
+        # Launch interface
+        interface.launch(share=True)
+    except Exception as e:
+        logger.error(f"Error launching app: {e}")
+        import traceback
+        logger.error(traceback.format_exc())
 
 if __name__ == "__main__":
     launch_app()