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()