Create app.py

app.py

@@ -0,0 +1,362 @@
+import os
+import gradio as gr
+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
+import cv2
+from io import BytesIO
+import urllib.request
+import tempfile
+import rasterio
+from rasterio.plot import reshape_as_image
+import warnings
+warnings.filterwarnings("ignore")
+
+# Set device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Define the DeepLabv3+ model architecture
+# This needs to match your trained model architecture
+from torchvision.models.segmentation import deeplabv3_resnet50
+
+# Initialize the model
+model = deeplabv3_resnet50(pretrained=False, num_classes=2)
+
+# Download model weights from HuggingFace
+MODEL_REPO = "dcrey7/wetlands_segmentation_deeplabsv3plus"
+MODEL_FILENAME = "DeepLabV3plus_best_model.pth"
+
+def download_model_weights():
+    """Download model weights from HuggingFace repository"""
+    try:
+        os.makedirs('weights', exist_ok=True)
+        local_path = os.path.join('weights', MODEL_FILENAME)
+        
+        # Check if weights are already downloaded
+        if os.path.exists(local_path):
+            print(f"Model weights already downloaded at {local_path}")
+            return local_path
+        
+        # Download weights
+        print(f"Downloading model weights from {MODEL_REPO}...")
+        url = f"https://huggingface.co/{MODEL_REPO}/resolve/main/{MODEL_FILENAME}"
+        urllib.request.urlretrieve(url, local_path)
+        print(f"Model weights downloaded to {local_path}")
+        return local_path
+    except Exception as e:
+        print(f"Error downloading model weights: {e}")
+        return None
+
+# Load the model weights
+weights_path = download_model_weights()
+if weights_path:
+    try:
+        model.load_state_dict(torch.load(weights_path, map_location=device))
+        print("Model weights loaded successfully")
+    except Exception as e:
+        print(f"Error loading model weights: {e}")
+else:
+    print("No weights available. Model will not produce valid predictions.")
+
+model.to(device)
+model.eval()
+
+def preprocess_image(image, target_size=(128, 128)):
+    """
+    Preprocess an image for inference
+    """
+    # Convert to numpy array if PIL image
+    if isinstance(image, Image.Image):
+        image = np.array(image)
+    
+    # Ensure RGB format
+    if len(image.shape) == 2:  # Grayscale
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+    elif image.shape[2] == 4:  # RGBA
+        image = image[:, :, :3]
+    
+    # Resize image
+    image_resized = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
+    
+    # Normalize image to [0, 1]
+    image_normalized = image_resized.astype(np.float32)
+    if image_normalized.max() > 0:
+        image_normalized = image_normalized / image_normalized.max()
+    
+    # Convert to tensor [C, H, W]
+    image_tensor = torch.from_numpy(image_normalized.transpose(2, 0, 1)).float().unsqueeze(0)
+    return image_tensor, image_resized
+
+def preprocess_tiff(tiff_path, target_size=(128, 128)):
+    """
+    Preprocess a TIFF file for inference
+    """
+    try:
+        with rasterio.open(tiff_path) as src:
+            # Read RGB bands if available
+            if src.count >= 3:
+                red = src.read(1)
+                green = src.read(2)
+                blue = src.read(3)
+                image = np.dstack((red, green, blue))
+            else:
+                # If less than 3 bands, read all available bands
+                bands = [src.read(i+1) for i in range(src.count)]
+                # If only one band, duplicate to create RGB
+                if len(bands) == 1:
+                    image = np.dstack((bands[0], bands[0], bands[0]))
+                else:
+                    # Use available bands and pad with zeros if needed
+                    while len(bands) < 3:
+                        bands.append(np.zeros_like(bands[0]))
+                    image = np.dstack(bands[:3])  # Use first 3 bands
+        
+        # Normalize image to [0, 1]
+        image = image.astype(np.float32)
+        if image.max() > 0:
+            image = image / image.max()
+        
+        # Create a display image
+        display_image = (image * 255).astype(np.uint8)
+        
+        # Resize image
+        image_resized = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
+        display_resized = cv2.resize(display_image, target_size, interpolation=cv2.INTER_LINEAR)
+        
+        # Convert to tensor [C, H, W]
+        image_tensor = torch.from_numpy(image_resized.transpose(2, 0, 1)).float().unsqueeze(0)
+        
+        return image_tensor, display_resized
+    except Exception as e:
+        print(f"Error processing TIFF: {e}")
+        return None, None
+
+def preprocess_mask(mask, target_size=(128, 128)):
+    """
+    Preprocess a ground truth mask
+    """
+    # Convert to numpy array if PIL image
+    if isinstance(mask, Image.Image):
+        mask = np.array(mask)
+    
+    # Convert to grayscale if needed
+    if len(mask.shape) == 3:
+        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
+    
+    # Resize mask
+    mask_resized = cv2.resize(mask, target_size, interpolation=cv2.INTER_NEAREST)
+    
+    # Binarize the mask (0: background, 1: wetland)
+    mask_binary = (mask_resized > 127).astype(np.uint8)
+    
+    return mask_binary
+
+def predict_segmentation(image_tensor):
+    """
+    Run inference on the model
+    """
+    try:
+        image_tensor = image_tensor.to(device)
+        
+        with torch.no_grad():
+            output = model(image_tensor)
+            
+            # Extract the output based on model type
+            if isinstance(output, dict):
+                output = output['out']
+            
+            # Get the predicted class (0: background, 1: wetland)
+            pred = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()
+        
+        return pred
+    except Exception as e:
+        print(f"Error during prediction: {e}")
+        return None
+
+def calculate_metrics(pred_mask, gt_mask):
+    """
+    Calculate evaluation metrics between prediction and ground truth
+    """
+    # Ensure binary masks
+    pred_binary = (pred_mask > 0).astype(np.uint8)
+    gt_binary = (gt_mask > 0).astype(np.uint8)
+    
+    # Calculate intersection and union
+    intersection = np.logical_and(pred_binary, gt_binary).sum()
+    union = np.logical_or(pred_binary, gt_binary).sum()
+    
+    # Calculate IoU
+    iou = intersection / union if union > 0 else 0
+    
+    # Calculate precision and recall
+    true_positive = intersection
+    false_positive = pred_binary.sum() - true_positive
+    false_negative = gt_binary.sum() - true_positive
+    
+    precision = true_positive / (true_positive + false_positive) if (true_positive + false_positive) > 0 else 0
+    recall = true_positive / (true_positive + false_negative) if (true_positive + false_negative) > 0 else 0
+    
+    # Calculate F1 score
+    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
+    
+    metrics = {
+        "IoU": float(iou),
+        "Precision": float(precision),
+        "Recall": float(recall),
+        "F1 Score": float(f1)
+    }
+    
+    return metrics
+
+def process_images(input_image=None, input_tiff=None, gt_mask=None):
+    """
+    Process input images and generate predictions
+    """
+    try:
+        # Check if we have input
+        if input_image is None and input_tiff is None:
+            return None, "Please upload an image or TIFF file."
+        
+        # Process the input image
+        if input_tiff is not None:
+            # Save uploaded TIFF to a temporary file
+            with tempfile.NamedTemporaryFile(suffix='.tif', delete=False) as temp_tiff:
+                temp_tiff_path = temp_tiff.name
+                temp_tiff.write(input_tiff)
+            
+            # Process TIFF file
+            image_tensor, display_image = preprocess_tiff(temp_tiff_path)
+            os.unlink(temp_tiff_path)  # Clean up temp file
+        else:
+            # Process regular image
+            image_tensor, display_image = preprocess_image(input_image)
+        
+        if image_tensor is None:
+            return None, "Failed to process the input image."
+        
+        # Get prediction
+        pred_mask = predict_segmentation(image_tensor)
+        if pred_mask is None:
+            return None, "Failed to generate prediction."
+        
+        # Process ground truth mask if provided
+        gt_mask_processed = 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()])
+        
+        # Create visualization
+        fig = plt.figure(figsize=(12, 6))
+        
+        if gt_mask_processed is not None:
+            # Show original, ground truth, and prediction
+            plt.subplot(1, 3, 1)
+            plt.imshow(display_image)
+            plt.title("Input Image")
+            plt.axis('off')
+            
+            plt.subplot(1, 3, 2)
+            plt.imshow(gt_mask_processed, cmap='binary')
+            plt.title("Ground Truth")
+            plt.axis('off')
+            
+            plt.subplot(1, 3, 3)
+            plt.imshow(pred_mask, cmap='binary')
+            plt.title("Prediction")
+            plt.axis('off')
+        else:
+            # Show original and prediction
+            plt.subplot(1, 2, 1)
+            plt.imshow(display_image)
+            plt.title("Input Image")
+            plt.axis('off')
+            
+            plt.subplot(1, 2, 2)
+            plt.imshow(pred_mask, cmap='binary')
+            plt.title("Predicted Wetlands")
+            plt.axis('off')
+        
+        # Calculate wetland percentage
+        wetland_percentage = np.mean(pred_mask) * 100
+        
+        # Add metrics info
+        result_text = f"Wetland Coverage: {wetland_percentage:.2f}%"
+        if metrics_text:
+            result_text += f"\n\nEvaluation Metrics:\n{metrics_text}"
+        
+        # Convert figure to image
+        buf = BytesIO()
+        plt.tight_layout()
+        plt.savefig(buf, format='png', dpi=150)
+        buf.seek(0)
+        result_image = Image.open(buf)
+        plt.close(fig)
+        
+        return result_image, result_text
+        
+    except Exception as e:
+        print(f"Error in processing: {e}")
+        return None, f"Error: {str(e)}"
+
+# Create Gradio interface
+with gr.Blocks(title="Wetlands Segmentation from Satellite Imagery") as demo:
+    gr.Markdown("# Wetlands Segmentation from Satellite Imagery")
+    gr.Markdown("Upload a satellite image or TIFF file to identify wetland areas. Optionally, you can also upload a ground truth mask for evaluation.")
+    
+    with gr.Row():
+        with gr.Column():
+            # Input options
+            gr.Markdown("### Input")
+            with gr.Tab("Upload Image"):
+                input_image = gr.Image(label="Upload Satellite Image", type="numpy")
+            
+            with gr.Tab("Upload TIFF"):
+                input_tiff = gr.File(label="Upload TIFF File", file_types=[".tif", ".tiff"])
+            
+            gt_mask = gr.Image(label="Ground Truth Mask (Optional)", type="numpy")
+            
+            process_btn = gr.Button("Analyze Image", variant="primary")
+        
+        with gr.Column():
+            # Output
+            gr.Markdown("### Results")
+            output_image = gr.Image(label="Segmentation Results", type="pil")
+            output_text = gr.Textbox(label="Statistics", lines=6)
+    
+    # Information about the model
+    gr.Markdown("### About this model")
+    gr.Markdown("""
+    This application uses a DeepLabv3+ model trained to segment wetland areas in satellite imagery.
+    
+    **Model Details:**
+    - Architecture: DeepLabv3+ with ResNet-50 backbone
+    - Input: RGB satellite imagery
+    - Output: Binary segmentation mask (Wetland vs Background)
+    - Resolution: 128×128 pixels
+    
+    **Tips for best results:**
+    - The model works best with RGB satellite imagery
+    - For optimal results, use images with similar characteristics to those used in training
+    - The model focuses on identifying wetland regions in natural landscapes
+    
+    **Repository:** [dcrey7/wetlands_segmentation_deeplabsv3plus](https://huggingface.co/dcrey7/wetlands_segmentation_deeplabsv3plus)
+    """)
+    
+    # Set up event handlers
+    process_btn.click(
+        fn=process_images,
+        inputs=[input_image, input_tiff, gt_mask],
+        outputs=[output_image, output_text]
+    )
+
+# Launch the app
+demo.launch()