app.py

import os
import gradio as gr
import torch
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import requests
import io
import matplotlib.colors as mcolors
import cv2
from io import BytesIO
import urllib.request
import tempfile
import rasterio
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")

# Try to import albumentations if needed for preprocessing
try:
    import albumentations as A
    albumentations_available = True
    print("Successfully imported albumentations")
except ImportError:
    albumentations_available = False
    print("Warning: albumentations not available, will try to install it")
    import subprocess
    try:
        subprocess.check_call([
            "pip", "install", "albumentations"
        ])
        import albumentations as A
        albumentations_available = True
        print("Successfully installed and imported albumentations")
    except:
        print("Failed to install albumentations, will use OpenCV for transforms")

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# Initialize the model
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 valid 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"
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:
        # Try to load with strict=False to allow for some parameter mismatches
        state_dict = torch.load(weights_path, map_location=device)
        # 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}")
else:
    print("No weights available. Model will not produce valid predictions.")

model.to(device)
model.eval()

def read_tiff_image(tiff_path):
    """
    Read a TIFF image using rasterio, focusing on RGB bands
    This matches your training data loading approach
    """
    try:
        # Read the image using rasterio (get RGB channels)
        with rasterio.open(tiff_path) as src:
            # Check if we have enough bands
            if src.count >= 3:
                red = src.read(1)
                green = src.read(2)
                blue = src.read(3)
                
                # Stack to create RGB image
                image = np.dstack((red, green, blue)).astype(np.float32)
                
                # Normalize to [0, 1]
                if image.max() > 0:
                    image = image / image.max()
                
                return image
            else:
                # If less than 3 bands, handle accordingly
                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
                if image.max() > 0:
                    image = image / image.max()
                
                return image
    except Exception as e:
        print(f"Error reading TIFF file: {e}")
        return None

def read_tiff_mask(mask_path):
    """
    Read a TIFF mask using rasterio
    This matches your training data loading approach
    """
    try:
        # Read mask
        with rasterio.open(mask_path) as src:
            mask = src.read(1).astype(np.uint8)
        return mask
    except Exception as e:
        print(f"Error reading mask file: {e}")
        return None

def preprocess_image(image, target_size=(128, 128)):
    """
    Preprocess an image for inference
    """
    # If image is already a numpy array, use it directly
    if isinstance(image, np.ndarray):
        # 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]
        
        # Make a copy for display
        display_image = image.copy()
        
        # Normalize to [0, 1] if needed
        if display_image.max() > 1.0:
            image = image.astype(np.float32) / 255.0
    
    # Convert PIL image to numpy
    elif 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]
        
        # Make a copy for display
        display_image = image.copy()
        
        # Normalize to [0, 1]
        image = image.astype(np.float32) / 255.0
    else:
        print(f"Unsupported image type: {type(image)}")
        return None, None
    
    # Resize image to the target size
    if albumentations_available:
        # Use albumentations to match training preprocessing
        aug = A.Compose([
            A.PadIfNeeded(min_height=target_size[0], min_width=target_size[1], 
                         border_mode=cv2.BORDER_CONSTANT, value=0),
            A.CenterCrop(height=target_size[0], width=target_size[1])
        ])
        augmented = aug(image=image)
        image_resized = augmented['image']
    else:
        # Fallback to OpenCV
        image_resized = cv2.resize(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_image

def extract_file_content(file_obj):
    """Extract content from the file object, handling different types"""
    try:
        if hasattr(file_obj, 'name') and isinstance(file_obj, str):
            # Handle Gradio's NamedString
            content = file_obj
            if os.path.exists(content):
                # It's a path
                with open(content, 'rb') as f:
                    return f.read()
            else:
                # It's content
                return content.encode('latin1')
        elif hasattr(file_obj, 'read'):
            # File-like object
            return file_obj.read()
        elif isinstance(file_obj, bytes):
            # Already bytes
            return file_obj
        elif isinstance(file_obj, str):
            # String path
            if os.path.exists(file_obj):
                with open(file_obj, 'rb') as f:
                    return f.read()
            else:
                return file_obj.encode('utf-8')
        else:
            print(f"Unsupported file object type: {type(file_obj)}")
            return None
    except Exception as e:
        print(f"Error extracting file content: {e}")
        return None

def process_uploaded_tiff(file_obj):
    """Process an uploaded TIFF file"""
    try:
        # Get file content
        file_content = extract_file_content(file_obj)
        if file_content is None:
            print("Failed to extract file content")
            return None, None
        
        # Save to a temporary file
        with tempfile.NamedTemporaryFile(delete=False, suffix='.tif') as temp_file:
            temp_path = temp_file.name
            temp_file.write(file_content)
        
        # Read as TIFF
        image = read_tiff_image(temp_path)
        
        # Clean up
        os.unlink(temp_path)
        
        if image is None:
            return None, None
        
        # Make a copy for display
        display_image = (image * 255).astype(np.uint8) if image.max() <= 1.0 else image.copy()
        
        # Resize/preprocess
        if albumentations_available:
            aug = A.Compose([
                A.PadIfNeeded(min_height=128, min_width=128, 
                             border_mode=cv2.BORDER_CONSTANT, value=0),
                A.CenterCrop(height=128, width=128)
            ])
            augmented = aug(image=image)
            image_resized = augmented['image']
        else:
            image_resized = cv2.resize(image, (128, 128), interpolation=cv2.INTER_LINEAR)
        
        # Convert to tensor
        image_tensor = torch.from_numpy(image_resized.transpose(2, 0, 1)).float().unsqueeze(0)
        
        return image_tensor, display_image
    
    except Exception as e:
        print(f"Error processing uploaded TIFF: {e}")
        import traceback
        traceback.print_exc()
        return None, None

def process_uploaded_mask(file_obj):
    """Process an uploaded mask file"""
    try:
        # Get file content
        file_content = extract_file_content(file_obj)
        if file_content is None:
            return None
        
        # Save to a temporary file
        # Determine suffix based on file name if available
        suffix = '.tif'
        if hasattr(file_obj, 'name'):
            file_name = getattr(file_obj, 'name')
            if isinstance(file_name, str) and '.' in file_name:
                suffix = '.' + file_name.split('.')[-1].lower()
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:
            temp_path = temp_file.name
            temp_file.write(file_content)
        
        # Check if it's a TIFF file
        if temp_path.lower().endswith(('.tif', '.tiff')):
            mask = read_tiff_mask(temp_path)
        else:
            # Try to open as a regular image
            try:
                mask_img = Image.open(temp_path)
                mask = np.array(mask_img)
                if len(mask.shape) == 3:
                    mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
            except Exception as e:
                print(f"Error opening mask as regular image: {e}")
                os.unlink(temp_path)
                return None
        
        # Clean up
        os.unlink(temp_path)
        
        if mask is None:
            return None
        
        # Resize mask to 128x128
        if albumentations_available:
            aug = A.Compose([
                A.PadIfNeeded(min_height=128, min_width=128, 
                             border_mode=cv2.BORDER_CONSTANT, value=0),
                A.CenterCrop(height=128, width=128)
            ])
            augmented = aug(image=mask)
            mask_resized = augmented['image']
        else:
            mask_resized = cv2.resize(mask, (128, 128), interpolation=cv2.INTER_NEAREST)
        
        # Binarize the mask (0: background, 1: wetland)
        mask_binary = (mask_resized > 0).astype(np.uint8)
        
        return mask_binary
    
    except Exception as e:
        print(f"Error processing uploaded mask: {e}")
        import traceback
        traceback.print_exc()
        return None

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)
            
            # Handle different model output formats
            if isinstance(output, dict):
                output = output['out']
            if output.shape[1] > 1:  # Multi-class output
                pred = torch.argmax(output, dim=1).squeeze(0).cpu().numpy()
            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 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_file=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
        if input_tiff is not None and input_tiff:
            # Process uploaded TIFF file
            image_tensor, display_image = process_uploaded_tiff(input_tiff)
            if image_tensor is None:
                return None, "Failed to process the input TIFF file."
        elif input_image is not None:
            # Process regular image
            image_tensor, display_image = preprocess_image(input_image)
            if image_tensor is None:
                return None, "Failed to process the input image."
        else:
            return None, "No valid input provided."
        
        # 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_file is not None and gt_mask_file:
            gt_mask_processed = process_uploaded_mask(gt_mask_file)
            
            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))
        
        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 for display
        plt.tight_layout()
        buf = BytesIO()
        plt.savefig(buf, format='png')
        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}")
        import traceback
        traceback.print_exc()
        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"])
            
            # Ground truth mask as file upload
            gt_mask_file = gr.File(label="Ground Truth Mask (Optional)", file_types=[".tif", ".tiff", ".png", ".jpg", ".jpeg"])
            
            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-34
    - 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
    - For ground truth masks, both TIFF and standard image formats are supported
    
    **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_file],
        outputs=[output_image, output_text]
    )

# Launch the app
demo.launch()