Update app.py

app.py

@@ -4,7 +4,6 @@ import torch
 import numpy as np
 import matplotlib.pyplot as plt
 from PIL import Image
-import torchvision.transforms as transforms
 import requests
 import io
 import matplotlib.colors as mcolors
@@ -17,27 +16,49 @@ from rasterio.plot import reshape_as_image
 import warnings
 warnings.filterwarnings("ignore")
 
+# Try to import segmentation_models_pytorch
+try:
+    import segmentation_models_pytorch as smp
+    smp_available = True
+    print("Successfully imported segmentation_models_pytorch")
+except ImportError:
+    smp_available = False
+    print("Warning: segmentation_models_pytorch not available, will try to install it")
+    import subprocess
+    try:
+        subprocess.check_call([
+            "pip", "install", "segmentation-models-pytorch"
+        ])
+        import segmentation_models_pytorch as smp
+        smp_available = True
+        print("Successfully installed and imported segmentation_models_pytorch")
+    except:
+        print("Failed to install segmentation_models_pytorch")
+
 # Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-# Define a custom DeepLabV3+ model that matches your trained model architecture
-class DeepLabV3Plus(torch.nn.Module):
-    def __init__(self, num_classes=2):
-        super(DeepLabV3Plus, self).__init__()
-        self.encoder = torch.nn.Sequential() # ResNet backbone
-        self.decoder = torch.nn.Sequential() # Decoder modules 
-        self.segmentation_head = torch.nn.Conv2d(256, num_classes, kernel_size=1)
-        
-    def forward(self, x):
-        # Forward pass (simplified since we're only using this for loading weights)
-        features = self.encoder(x)
-        decoder_output = self.decoder(features)
-        masks = self.segmentation_head(decoder_output)
-        return masks
-
 # Initialize the model
-model = DeepLabV3Plus(num_classes=2)
+if smp_available:
+    # Define the DeepLabV3+ model using smp
+    model = smp.DeepLabV3Plus(
+        encoder_name="resnet34",  # Using ResNet34 backbone as in your training
+        encoder_weights=None,     # We'll load your custom weights
+        in_channels=3,            # RGB input
+        classes=1,                # Binary segmentation
+    )
+else:
+    # Fallback to a simple model that won't actually work but allows the UI to load
+    print("Warning: Using a placeholder model that won't produce correct predictions.")
+    from torch import nn
+    class PlaceholderModel(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.conv = nn.Conv2d(3, 1, 3, padding=1)
+        def forward(self, x):
+            return self.conv(x)
+    model = PlaceholderModel()
 
 # Download model weights from HuggingFace
 MODEL_REPO = "dcrey7/wetlands_segmentation_deeplabsv3plus"
@@ -64,31 +85,25 @@ def download_model_weights():
         print(f"Error downloading model weights: {e}")
         return None
 
-# Dummy model for testing if model weights can't be loaded
-class DummyModel(torch.nn.Module):
-    def __init__(self):
-        super(DummyModel, self).__init__()
-        
-    def forward(self, x):
-        # Simply return a random segmentation mask for visualization
-        batch_size, _, height, width = x.shape
-        return torch.randint(0, 2, (batch_size, 2, height, width), device=x.device).float()
-
 # Load the model weights
 weights_path = download_model_weights()
 if weights_path:
     try:
-        # Try to load using strict=False to allow for partial weight loading
+        # Try to load with strict=False to allow for some parameter mismatches
         state_dict = torch.load(weights_path, map_location=device)
-        model.load_state_dict(state_dict, strict=False)
-        print("Model weights loaded with non-strict mapping")
+        # Check if we need to modify the state dict keys
+        if all(key.startswith('encoder.') or key.startswith('decoder.') for key in list(state_dict.keys())[:5]):
+            print("Model weights use encoder/decoder format, loading directly")
+            model.load_state_dict(state_dict, strict=False)
+        else:
+            print("Attempting to adapt state dict to match model architecture")
+            # This is a placeholder for state dict adaptation if needed
+            model.load_state_dict(state_dict, strict=False)
+        print("Model weights loaded successfully")
     except Exception as e:
         print(f"Error loading model weights: {e}")
-        print("Using dummy model for demo purposes")
-        model = DummyModel()
 else:
-    print("No weights available. Using dummy model.")
-    model = DummyModel()
+    print("No weights available. Model will not produce valid predictions.")
 
 model.to(device)
 model.eval()
@@ -167,8 +182,22 @@ def preprocess_mask(mask, target_size=(128, 128)):
     """
     Preprocess a ground truth mask
     """
+    # If mask is a file path (string), open it
+    if isinstance(mask, str):
+        try:
+            # Try to open as a TIFF file with rasterio
+            with rasterio.open(mask) as src:
+                mask_array = src.read(1)  # Read first band
+                mask = mask_array
+        except:
+            # Fall back to opening with PIL
+            try:
+                mask = np.array(Image.open(mask))
+            except Exception as e:
+                print(f"Error reading mask file: {e}")
+                return None
     # Convert to numpy array if PIL image
-    if isinstance(mask, Image.Image):
+    elif isinstance(mask, Image.Image):
         mask = np.array(mask)
     
     # Convert to grayscale if needed
@@ -193,22 +222,18 @@ def predict_segmentation(image_tensor):
         with torch.no_grad():
             output = model(image_tensor)
             
-            # Get the predicted class (0: background, 1: wetland)
-            # Handle different output formats
-            if isinstance(output, dict) and 'out' in output:
+            # Handle different model output formats
+            if isinstance(output, dict):
                 output = output['out']
-            
-            if output.shape[1] > 1:  # If output has multiple channels (classes)
+            if output.shape[1] > 1:  # Multi-class output
                 pred = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()
-            else:
-                # If output is single channel, threshold it
-                pred = (output.squeeze(0).squeeze(0) > 0.5).cpu().numpy().astype(np.uint8)
-        
+            else:  # Binary output (from smp models)
+                pred = (torch.sigmoid(output) > 0.5).squeeze().cpu().numpy().astype(np.uint8)
+                
         return pred
     except Exception as e:
         print(f"Error during prediction: {e}")
-        # Return random prediction for demo purposes
-        return np.random.randint(0, 2, (128, 128), dtype=np.uint8)
+        return None
 
 def calculate_metrics(pred_mask, gt_mask):
     """
@@ -245,29 +270,6 @@ def calculate_metrics(pred_mask, gt_mask):
     
     return metrics
 
-def save_uploaded_file(file_obj):
-    """Save an uploaded file to a temporary location and return the path"""
-    try:
-        # Create a temporary file
-        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.tif')
-        temp_path = temp_file.name
-        
-        # Write the content to the file
-        if hasattr(file_obj, 'name'):
-            # If it's a FileUpload object from gradio
-            with open(file_obj.name, 'rb') as f:
-                content = f.read()
-                temp_file.write(content)
-        else:
-            # If it's binary content
-            temp_file.write(file_obj)
-            
-        temp_file.close()
-        return temp_path
-    except Exception as e:
-        print(f"Error saving uploaded file: {e}")
-        return None
-
 def process_images(input_image=None, input_tiff=None, gt_mask=None):
     """
     Process input images and generate predictions
@@ -279,10 +281,18 @@ def process_images(input_image=None, input_tiff=None, gt_mask=None):
         
         # Process the input image
         if input_tiff is not None:
-            # Save uploaded TIFF to a temporary file
-            temp_tiff_path = save_uploaded_file(input_tiff)
-            if not temp_tiff_path:
-                return None, "Failed to process the uploaded TIFF file."
+            # Create a temporary file for the uploaded TIFF
+            with tempfile.NamedTemporaryFile(suffix='.tif', delete=False) as temp_tiff:
+                temp_tiff_path = temp_tiff.name
+                
+                # Write the file content to the temporary file
+                if isinstance(input_tiff, str):
+                    # If input_tiff is a path
+                    with open(input_tiff, 'rb') as f:
+                        temp_tiff.write(f.read())
+                else:
+                    # If input_tiff is file-like object or bytes
+                    temp_tiff.write(input_tiff)
             
             # Process TIFF file
             image_tensor, display_image = preprocess_tiff(temp_tiff_path)
@@ -308,9 +318,28 @@ def process_images(input_image=None, input_tiff=None, gt_mask=None):
         metrics_text = ""
         
         if gt_mask is not None:
-            gt_mask_processed = preprocess_mask(gt_mask)
-            metrics = calculate_metrics(pred_mask, gt_mask_processed)
-            metrics_text = "\n".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
+            # If gt_mask is a file upload
+            if isinstance(gt_mask, (str, bytes)):
+                # Create a temporary file for the mask
+                with tempfile.NamedTemporaryFile(suffix='.tif', delete=False) as temp_mask:
+                    temp_mask_path = temp_mask.name
+                    if isinstance(gt_mask, str):
+                        with open(gt_mask, 'rb') as f:
+                            temp_mask.write(f.read())
+                    else:
+                        temp_mask.write(gt_mask)
+                gt_mask_processed = preprocess_mask(temp_mask_path)
+                try:
+                    os.unlink(temp_mask_path)
+                except:
+                    pass
+            else:
+                # Normal image upload
+                gt_mask_processed = preprocess_mask(gt_mask)
+            
+            if gt_mask_processed is not None:
+                metrics = calculate_metrics(pred_mask, gt_mask_processed)
+                metrics_text = "\n".join([f"{k}: {v:.4f}" for k, v in metrics.items()])
         
         # Create visualization
         fig = plt.figure(figsize=(12, 6))
@@ -362,10 +391,9 @@ def process_images(input_image=None, input_tiff=None, gt_mask=None):
         return result_image, result_text
         
     except Exception as e:
-        import traceback
-        trace = traceback.format_exc()
         print(f"Error in processing: {e}")
-        print(trace)
+        import traceback
+        traceback.print_exc()
         return None, f"Error: {str(e)}"
 
 # Create Gradio interface
@@ -399,7 +427,7 @@ with gr.Blocks(title="Wetlands Segmentation from Satellite Imagery") as demo:
     This application uses a DeepLabv3+ model trained to segment wetland areas in satellite imagery.
     
     **Model Details:**
-    - Architecture: DeepLabv3+ with ResNet-50 backbone
+    - Architecture: DeepLabv3+ with ResNet-34
     - Input: RGB satellite imagery
     - Output: Binary segmentation mask (Wetland vs Background)
     - Resolution: 128×128 pixels