dcrey7's picture
Update app.py
6656678 verified
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
import pandas as pd
import joblib
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 segmentation 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 segmentation model weights from HuggingFace
SEGMENTATION_MODEL_REPO = "dcrey7/wetlands_segmentation_deeplabsv3plus"
SEGMENTATION_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', SEGMENTATION_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 {SEGMENTATION_MODEL_REPO}...")
url = f"https://huggingface.co/{SEGMENTATION_MODEL_REPO}/resolve/main/{SEGMENTATION_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 segmentation 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()
# Load the cloud detection model
def load_cloud_detection_model():
"""Load cloud detection model from the local file"""
try:
# Check if the model file exists
model_path = "cloud_detection_lightgbm.joblib"
if os.path.exists(model_path):
# Load the model
cloud_model = joblib.load(model_path)
print(f"Cloud detection model loaded successfully from {model_path}")
return cloud_model
else:
print(f"Cloud detection model file not found at {model_path}")
return None
except Exception as e:
print(f"Error loading cloud detection model: {e}")
return None
# Load the cloud detection model
cloud_model = load_cloud_detection_model()
if cloud_model:
print("Cloud detection model is ready for predictions")
else:
print("Warning: Cloud detection model could not be loaded")
def normalize(band):
"""Normalize band values using 2-98 percentile range"""
# Handle potential NaN or inf values
band_cleaned = band[np.isfinite(band)]
if len(band_cleaned) == 0:
return band
# Use percentiles to avoid outliers
band_min, band_max = np.percentile(band_cleaned, (2, 98))
# Avoid division by zero
if band_max == band_min:
return np.zeros_like(band)
band_normalized = (band - band_min) / (band_max - band_min)
band_normalized = np.clip(band_normalized, 0, 1)
return band_normalized
def calculate_cv(band):
"""Calculate coefficient of variation (CV) for a band"""
# First normalize the band
band_normalized = normalize(band)
# Handle potential NaN or inf values
band_cleaned = band_normalized[np.isfinite(band_normalized)]
if len(band_cleaned) == 0:
return 0
# Get mean and std dev
mean = np.mean(band_cleaned)
# Guard against division by zero or very small means
if abs(mean) < 1e-10:
return 0
std = np.std(band_cleaned)
cv = (std / mean) # CV as ratio (not percentage)
return cv
def read_tiff_image_for_segmentation(tiff_path):
"""
Read a TIFF image using rasterio, focusing on RGB bands (first 3 bands)
for wetland segmentation
"""
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 for segmentation: {e}")
return None
def extract_cloud_features_from_tiff(tiff_path):
"""
Extract CV features from all bands in a TIFF file for cloud detection.
Will try to use up to 10 bands.
"""
try:
with rasterio.open(tiff_path) as src:
num_bands = min(src.count, 10) # Use up to 10 bands
# Process each band
features = {}
for i in range(1, num_bands + 1):
band = src.read(i)
# Calculate coefficient of variation
cv_value = calculate_cv(band)
# Store feature with name matching the training data
features[f'band{i}_cv'] = cv_value
# If we have fewer than 10 bands, fill the missing ones with zeros
for i in range(num_bands + 1, 11):
features[f'band{i}_cv'] = 0.0
return features
except Exception as e:
print(f"Error extracting cloud features from TIFF: {e}")
import traceback
traceback.print_exc()
return None
def extract_cloud_features_from_rgb(image):
"""
Extract CV features from RGB image for cloud detection.
Will use 3 bands and fill the remaining 7 with zeros to match the expected 10 features.
"""
try:
# Make sure image is in float format in range [0,1]
if image.dtype != np.float32 and image.dtype != np.float64:
image = image.astype(np.float32)
if image.max() > 1.0:
image = image / 255.0
# Create a dictionary for band CV features
features = {}
# Process each channel/band
for i in range(min(1, image.shape[2])):
band = image[:, :, i]
cv_value = calculate_cv(band)
features[f'band{i+1}_cv'] = cv_value
# Fill remaining bands with zeros to match the expected 10 features
# for i in range(4, 11):
# features[f'band{i}_cv'] = 0.0
return features
except Exception as e:
print(f"Error extracting cloud features from RGB: {e}")
import traceback
traceback.print_exc()
return None
def predict_cloud(features_dict, model):
"""Predict if an image is cloudy based on extracted features"""
if model is None:
return {'prediction': 'Model unavailable', 'probability': 0.0}
try:
# Ensure all 10 features from band1_cv to band10_cv are present
feature_dict = {}
for i in range(1, 11):
feature_name = f'band{i}_cv'
feature_dict[feature_name] = features_dict.get(feature_name, 0.0)
# Create a DataFrame with all required features
feature_df = pd.DataFrame([feature_dict])
# Enable shape check disabling for prediction
if hasattr(model, 'set_params'):
model.set_params(predict_disable_shape_check=True)
# Make prediction
if hasattr(model, 'predict_proba'):
proba = model.predict_proba(feature_df)
if proba.shape[1] > 1: # Binary classification with probabilities for both classes
probability = proba[0][1] # Probability of the positive class (cloudy)
else:
probability = proba[0][0] # Single probability output
else:
# If model doesn't have predict_proba, use predict and assume binary output
pred = model.predict(feature_df)
probability = float(pred[0])
# Classification based on probability threshold
prediction = 'Cloudy' if probability >= 0.5 else 'Non-Cloudy'
return {
'prediction': prediction,
'probability': probability
}
except Exception as e:
print(f"Error predicting cloud: {e}")
import traceback
traceback.print_exc()
return {'prediction': 'Error', 'probability': 0.0}
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 for both segmentation and cloud detection"""
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, 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 for segmentation (only using first 3 bands)
image_for_segmentation = read_tiff_image_for_segmentation(temp_path)
# Extract cloud features from all available bands
cloud_features = extract_cloud_features_from_tiff(temp_path)
# Clean up
os.unlink(temp_path)
if image_for_segmentation is None:
return None, None, None
# Make a copy for display
display_image = (image_for_segmentation * 255).astype(np.uint8) if image_for_segmentation.max() <= 1.0 else image_for_segmentation.copy()
# Resize/preprocess for segmentation model
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_for_segmentation)
image_resized = augmented['image']
else:
image_resized = cv2.resize(image_for_segmentation, (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, cloud_features
except Exception as e:
print(f"Error processing uploaded TIFF: {e}")
import traceback
traceback.print_exc()
return None, 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, generate predictions for both wetland segmentation and cloud detection
"""
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 and initialize cloud_features
cloud_features = None
if input_tiff is not None and input_tiff:
# Process uploaded TIFF file for both segmentation and cloud detection
image_tensor, display_image, cloud_features = 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."
# For RGB images, we need to extract cloud features separately
cloud_features = extract_cloud_features_from_rgb(display_image)
else:
return None, "No valid input provided."
# Get wetland segmentation prediction
pred_mask = predict_segmentation(image_tensor)
if pred_mask is None:
return None, "Failed to generate wetland segmentation prediction."
# Get cloud prediction
cloud_result = {'prediction': 'Unknown', 'probability': 0.0}
if cloud_features and cloud_model:
cloud_result = predict_cloud(cloud_features, cloud_model)
# 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 results information
result_text = f"Wetland Coverage: {wetland_percentage:.2f}%\n\n"
# Add cloud detection results
result_text += f"Cloud Detection: {cloud_result['prediction']} "
result_text += f"({cloud_result['probability']*100:.2f}% confidence)\n\n"
# Add segmentation metrics if available
if metrics_text:
result_text += f"Evaluation 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 & Cloud Detection") as demo:
gr.Markdown("# Wetlands Segmentation & Cloud Detection from Satellite Imagery")
gr.Markdown("Upload a satellite image or TIFF file to identify wetland areas and detect cloud cover. 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=8)
# Information about the models
gr.Markdown("### About these models")
gr.Markdown("""
This application uses two models:
**1. Wetland Segmentation Model:**
- Architecture: DeepLabv3+ with ResNet-34
- Input: RGB satellite imagery
- Output: Binary segmentation mask (Wetland vs Background)
- Resolution: 128×128 pixels
**2. Cloud Detection Model:**
- Architecture: LightGBM Classifier
- Input: CV features extracted from up to 10 image bands
- Output: Binary classification (Cloudy vs Non-Cloudy) with probability
**Tips for best results:**
- For Cloudy image - train_11202327_p1, for Non cloudy image - train_02202325_p1
- For Cloudy image - test_07202330_p1, for Non cloudy image - test_02202325_p1
- The models work best with multi-band satellite imagery (TIFF files)
- For optimal cloud detection results, use TIFF files with 10 bands
- For optimal results, use images with similar characteristics to those used in training
- The wetland model focuses on identifying wetland regions in natural landscapes
- The cloud model detects cloud cover based on image band statistics
- 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()