Segmentation workflow using PyTorch (U-Net/CNN) for flood and water detection

!pip install -q torch torchvision gdown

import os
import glob
import math
from datetime import datetime, timedelta
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset, random_split
import torchvision.transforms as transforms
import torchvision
import pandas as pd
import numpy as np
import geopandas as gpd
import rasterio
from rasterio.features import shapes
import tqdm
import xarray as xr
import rioxarray as rxr
import planetary_computer
import pystac_client
import odc.stac
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from shapely.geometry import box, shape
from sklearn.metrics import jaccard_score, accuracy_score
from skimage.transform import resize
import dask
import dask.array as da
import dask.delayed

# Download the data (subset_vbos_smaller.zip)
!gdown "https://drive.google.com/uc?id=1rLNwGQf8C93CXpdQ293zKaKiQejDdJnp"

# Unzip the compressed data
!mkdir -p vbos/flood/
!unzip subset_vbos_smaller.zip -d vbos/flood/

# Get list of all labels
mask_paths = glob.glob('vbos/flood/subset_vbos_smaller/tst-ground-truth-flood-masks/*.tif')

len(mask_paths)

ls vbos/flood/subset_vbos_smaller/tst-ground-truth-flood-masks/ | head

cd vbos/flood/subset_vbos_smaller/

# Define U-Net model with ResNet34 backbone
class UNetWithResNetEncoder(nn.Module):
    def __init__(self, in_channels=2, out_channels=4):
        super().__init__()

        # Load pre-trained ResNet34 backbone
        resnet = torchvision.models.resnet34(weights=torchvision.models.resnet.ResNet34_Weights.DEFAULT)

        # Modify first convolution to accept 2 channels
        self.encoder1 = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False),
            resnet.bn1,
            resnet.relu
        )

        # Encoder layers
        self.encoder2 = resnet.layer1  # (B, 64, H/4, W/4)
        self.encoder3 = resnet.layer2  # (B, 128, H/8, W/8)
        self.encoder4 = resnet.layer3  # (B, 256, H/16, W/16)
        self.encoder5 = resnet.layer4  # (B, 512, H/32, W/32)

        # Bottleneck
        self.bottleneck = nn.Conv2d(512, 1024, kernel_size=3, padding=1)

        # Decoder with skip connections
        self.upconv1 = self._upsample(1024, 512)
        self.upconv2 = self._upsample(512 + 512, 256)  # Add encoder5 skip connection
        self.upconv3 = self._upsample(256 + 256, 128)  # Add encoder4 skip connection
        self.upconv4 = self._upsample(128 + 128, 64)   # Add encoder3 skip connection
        self.upconv5 = self._upsample(64 + 64, 64)     # Add encoder2 skip connection

        # Final segmentation output
        self.final_conv = nn.Conv2d(64, out_channels, kernel_size=1)

    def _upsample(self, in_channels, out_channels):
        """Helper function to create upsampling layers."""
        return nn.Sequential(
            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        # Encoder path
        enc1 = self.encoder1(x)   # (B, 64, 112, 112)
        enc2 = self.encoder2(enc1) # (B, 64, 56, 56)
        enc3 = self.encoder3(enc2) # (B, 128, 28, 28)
        enc4 = self.encoder4(enc3) # (B, 256, 14, 14)
        enc5 = self.encoder5(enc4) # (B, 512, 7, 7)

        # Bottleneck
        x = self.bottleneck(enc5)  # (B, 1024, 7, 7)
        x = self.upconv1(x)        # (B, 512, 14, 14)

        # Ensure enc5 matches x before concatenation
        enc5 = F.interpolate(enc5, size=x.shape[2:], mode="bilinear", align_corners=True)
        x = torch.cat([x, enc5], dim=1)

        x = self.upconv2(x)  # (B, 256, 28, 28)

        # Ensure enc4 matches x before concatenation
        enc4 = F.interpolate(enc4, size=x.shape[2:], mode="bilinear", align_corners=True)
        x = torch.cat([x, enc4], dim=1)

        x = self.upconv3(x)  # (B, 128, 56, 56)

        # Ensure enc3 matches x before concatenation
        enc3 = F.interpolate(enc3, size=x.shape[2:], mode="bilinear", align_corners=True)
        x = torch.cat([x, enc3], dim=1)

        x = self.upconv4(x)  # (B, 64, 112, 112)

        # Ensure enc2 matches x before concatenation
        enc2 = F.interpolate(enc2, size=x.shape[2:], mode="bilinear", align_corners=True)
        x = torch.cat([x, enc2], dim=1)

        x = self.upconv5(x)  # (B, 64, 224, 224)

        # Ensure final output size is exactly (B, 4, 224, 224)
        x = self.final_conv(x)  # (B, 4, 224, 224)
        x = F.interpolate(x, size=(224, 224), mode="bilinear", align_corners=True)  # Fix output size

        return x

# Function to fetch Sentinel-1 RTC data matching a given date (or up to 3 days after)
def fetch_s1_rtc_for_date(bounds, crs, date):
    catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")

    # Define the search time range (3-day + window)
    date = datetime.strptime(date, "%Y-%m-%d")  # Convert to datetime object

    start_date = date.strftime("%Y-%m-%d")
    end_date = (date + timedelta(days=3)).strftime("%Y-%m-%d")

    search = catalog.search(
        collections=["sentinel-1-rtc"],
        bbox=[bounds.left, bounds.bottom, bounds.right, bounds.top],
        datetime=f"{start_date}/{end_date}"
    )

    items = list(search.items())
    if not items:
        return None  # No matching data found

    signed_items = [planetary_computer.sign(item) for item in items]

    # Load the first available scene (or refine selection criteria as needed)
    ds = odc.stac.load(signed_items[:1], bbox=bounds, dtype="float32")
    ds = ds.rio.reproject(dst_crs=crs)  # Reproject to same CRS

    return ds.isel(time=0) #.mean(dim="time") if ds else None  # Take mean if multiple images exist

# Function to find matching Sentinel-1 RTC imagery for each row in DataFrame
def match_s1_rtc_to_dataframe(df, mask_path):
    filtered_row = df.loc[df["File Name"] == mask_path].iloc[0]
    sat_date = filtered_row["Satellite Date"]
    formatted_date = datetime.strptime(sat_date, "%d/%m/%Y").strftime("%Y-%m-%d")


    with rasterio.open(f"tst-ground-truth-flood-masks/{mask_path}_ground_truth.tif") as src:
        bounds = src.bounds  # Get bounding box of the mask
        crs = src.crs  # Coordinate Reference System

    s1_image = fetch_s1_rtc_for_date(bounds, crs, formatted_date)

    return s1_image

df = pd.read_csv('metadata/metadata.csv')
# Filter for specific countries
selected_countries = {"Vanuatu", "Tonga"} #, "Timor-Leste" , "Philippines", "Viet Nam", "Australia"}
df = df[df["Country"].isin(selected_countries)]

"""
# Limit the number of rows for the Philippines
df_philippines = df[df["Country"] == "Philippines"].head(5)  # Keep only 10 rows
df_other = df[df["Country"] != "Philippines"]  # Keep all other countries

# Concatenate back
df = pd.concat([df_philippines, df_other])

# Reset index
df = df.reset_index(drop=True)
"""

print("len(df):", len(df))

class FloodDataset(Dataset):
    def __init__(self, mask_paths, transform=None, tile_size=224):
        #self.mask_paths = filter_masks_within_vanuatu(mask_paths)
        self.mask_paths = list(df["File Name"])
        self.transform = transform
        self.image_list = []
        self.mask_list = []
        self.tile_size = tile_size

        # Use Dask Delayed for parallel fetching
        delayed_images = []
        for mask_path in self.mask_paths:
            s1_image = match_s1_rtc_to_dataframe(df, mask_path)
            delayed_images.append(dask.delayed(s1_image))

        # Compute all Dask delayed tasks at once (efficient batch processing)
        computed_images = dask.compute(*delayed_images)

        self.image_tiles = []
        self.mask_tiles = []

        for i, mask_path in enumerate(self.mask_paths):
            if computed_images[i] is None:
                continue  # Skip if no data was found

            with rasterio.open(f"tst-ground-truth-flood-masks/{mask_path}_ground_truth.tif") as src:
                mask = src.read(1).astype(np.float32)
            try:
                image_tensor = torch.tensor(computed_images[i].to_array().values, dtype=torch.float32)
                mask_tensor = torch.tensor(mask).unsqueeze(0)

                image_tiles, mask_tiles = self._tile_image_and_mask(image_tensor, mask_tensor)
                self.image_tiles.extend(image_tiles)  # Flatten the dataset
                self.mask_tiles.extend(mask_tiles)
            except:
                print("error when creating tensor")
                continue


    def __len__(self):
        return len(self.image_tiles)

    def __getitem__(self, idx):
        return self.image_tiles[idx], self.mask_tiles[idx]

    def _tile_image_and_mask(self, image, mask):
        #Tiles both the image and mask into patches of size (tile_size x tile_size).

        height, width = image.shape[-2:]  # Ensure compatibility with tensor dimensions
        image_tiles = []
        mask_tiles = []

        for i in range(0, height, self.tile_size):
            for j in range(0, width, self.tile_size):
                # Extract tiles
                image_tile = image[..., i:i + self.tile_size, j:j + self.tile_size]
                mask_tile = mask[..., i:i + self.tile_size, j:j + self.tile_size]

                if mask_tile.shape[-2] == self.tile_size and mask_tile.shape[-1] == self.tile_size and image_tile.shape[-2] == self.tile_size and image_tile.shape[-1] == self.tile_size:
                    image_tiles.append(image_tile.clone().detach())
                    mask_tiles.append(mask_tile.clone().detach())

        return image_tiles, mask_tiles

full_dataset = FloodDataset(mask_paths=mask_paths,
                             transform=transforms.Compose([transforms.Normalize(0, 1)]), tile_size=224)

# Define split sizes
train_size = int(0.7 * len(full_dataset))
val_size = int(0.20 * len(full_dataset))
test_size = len(full_dataset) - train_size - val_size  # Ensures total sums up correctly

# Perform split
generator = torch.Generator().manual_seed(42)
train_dataset, val_dataset, test_dataset = random_split(full_dataset, [train_size, val_size, test_size], generator)
print(f"Dataset sizes -> Train: {len(train_dataset)}, Val: {len(val_dataset)}, Test: {len(test_dataset)}")

# Create DataLoaders
train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, drop_last=True)
val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False, drop_last=True)
test_loader = DataLoader(test_dataset, batch_size=2, shuffle=False, drop_last=True)

epochs = 20

# Model, Loss, Optimizer
model = UNetWithResNetEncoder()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

print("Starting training")

# Training Loop
for epoch in tqdm.trange(epochs):
    model.train()  # Ensure the model is in training mode

    # Training phase
    train_loss = 0.0
    for batch in train_loader:
        images, masks = batch  # Unpack stacked tensors

        if isinstance(images, list):
            images = torch.cat(images, dim=0)  # Ensure 4D tensor (B, C, H, W)
        if isinstance(masks, list):
            masks = torch.cat(masks, dim=0)

        images = torch.stack(images) if isinstance(images, list) else images
        masks = torch.stack(masks) if isinstance(masks, list) else masks
        masks = masks.squeeze()
        masks = masks.long()  # Convert masks to LongTensor before computing loss

        optimizer.zero_grad()
        outputs = model(images)  # Forward pass
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()

        train_loss += loss.item()

    # Average training loss for the epoch
    train_loss /= len(train_loader)

    # Validation phase
    model.eval()  # Set the model to evaluation mode
    val_loss = 0.0
    with torch.no_grad():
        for batch in val_loader:
            images, masks = batch  # Unpack stacked tensors

            if isinstance(images, list):
                images = torch.cat(images, dim=0)
            if isinstance(masks, list):
                masks = torch.cat(masks, dim=0)

            images = torch.stack(images) if isinstance(images, list) else images
            masks = torch.stack(masks) if isinstance(masks, list) else masks
            masks = masks.squeeze()
            masks = masks.long()  # Convert masks to LongTensor

            outputs = model(images)  # Forward pass
            loss = criterion(outputs, masks)
            val_loss += loss.item()

    # Average validation loss for the epoch
    val_loss /= len(val_loader)

    print(f"Epoch {epoch+1}, Train Loss: {train_loss:.4f}, Validation Loss: {val_loss:.4f}")

# Directory to save the model checkpoints
MODEL_SAVE_DIR = '.'
os.makedirs(MODEL_SAVE_DIR, exist_ok=True)

# Save the final model after training
final_model_path = os.path.join(MODEL_SAVE_DIR, f"final_model_{epochs}ep.pth")
torch.save(model.state_dict(), final_model_path)
print(f"Final model saved at {final_model_path}")

# Inference Function
def inference(model, image):
    model.eval()
    with torch.no_grad():
        output = model(image.unsqueeze(0)) # 4 channels of logits
        prediction = torch.softmax(output, dim=1).squeeze(0)  # Softmax across channels
        prediction = torch.argmax(prediction, dim=0).cpu().numpy()  # Get class labels
    return prediction # 1, h, w

# Evaluation Function
def evaluate(model, data_loader):
    model.eval()
    iou_scores, accuracy_scores = [], []
    with torch.no_grad():
        for images, masks in data_loader:
            outputs = model(images)
            preds = torch.argmax(outputs, dim=1)  # Convert [batch, 4, H, W] → [batch, H, W]
            preds = preds.cpu().numpy().flatten()
            masks = masks.cpu().numpy().flatten()
            iou_scores.append(jaccard_score(masks, preds, average="macro"))
            accuracy_scores.append(accuracy_score(masks, preds))
    print(f"Mean IoU: {np.mean(iou_scores):.4f}, Mean Accuracy: {np.mean(accuracy_scores):.4f}")


# Run Evaluation
evaluate(model, test_loader)

# Define class labels and colormap
class_labels = {
    0: "No Data",
    1: "Land (not flooded)",
    2: "Flooded areas",
    3: "Permanent water bodies"
}

# Define discrete colormap
cmap = mcolors.ListedColormap(["black", "green", "blue", "navy"])  # Colors for classes 0-3
bounds = [0, 1, 2, 3, 4]  # Boundaries between color levels
norm = mcolors.BoundaryNorm(bounds, cmap.N)

def visualize_prediction(image, mask, prediction):
    fig, ax = plt.subplots(1, 3, figsize=(15, 5))

    # Convert 2-channel image to single-channel (VV)
    vv_image = image[0].cpu().numpy()  # Use first channel (VV)

    ax[0].imshow(vv_image, cmap="gray")  # Display as grayscale
    ax[0].set_title("Input Image (VV)")

    im1 = ax[1].imshow(mask.squeeze(), cmap=cmap, norm=norm)
    ax[1].set_title("Ground Truth")

    im2 = ax[2].imshow(prediction, cmap=cmap, norm=norm)
    ax[2].set_title("Prediction")

    # Create colorbar with class labels
    cbar = fig.colorbar(im2, ax=ax.ravel().tolist(), ticks=[0.5, 1.5, 2.5, 3.5], orientation="horizontal")
    cbar.ax.set_xticklabels([class_labels[i] for i in range(4)])  # Assign class names

    plt.show()

# Test Inference and Visualization
sample_image, sample_mask = test_dataset[0]
predicted_mask = inference(model, sample_image)
visualize_prediction(sample_image, sample_mask, predicted_mask)

admin_boundaries_gdf = gpd.read_file("./2016_phc_vut_iid_4326.geojson").to_crs(epsg="4326")

admin_boundaries_gdf

# Define the model architecture (must match the saved model)
model = UNetWithResNetEncoder(in_channels=2, out_channels=4)  # Ensure the same architecture

epochs = 20
# Load the model weights
model.load_state_dict(torch.load(f"final_model_{epochs}ep.pth"))

# Set the model to evaluation mode (important for inference)
model.eval()

#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#model.to(device)

# Function to fetch Sentinel-1 RTC data matching a given date (or up to 10 days after)
def fetch_s1_rtc_for_date_inference(bounds, date):
    catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")

    # Define the search time range (10-day + window)
    date = datetime.strptime(date, "%Y-%m-%d")  # Convert to datetime object

    start_date = date.strftime("%Y-%m-%d")
    end_date = (date + timedelta(days=15)).strftime("%Y-%m-%d")

    search = catalog.search(
        collections=["sentinel-1-rtc"],
        bbox=bounds,
        datetime=f"{start_date}/{end_date}"
    )

    items = list(search.items())
    if not items:
        return None  # No matching data found

    signed_items = [planetary_computer.sign(item) for item in items]

    # Load the first available scene (or refine selection criteria as needed)
    ds = odc.stac.load(signed_items, bbox=bounds, dtype="float32") #, chunks={'x': 512, 'y': 512}) # signed_items[:1]

    return ds.mean(dim="time") if ds else None  # Take mean if multiple images exist

def fetch_s2_rgb_for_date(bounds, date):
    catalog = pystac_client.Client.open("https://planetarycomputer.microsoft.com/api/stac/v1")

    date = datetime.strptime(date, "%Y-%m-%d")
    start_date = date.strftime("%Y-%m-%d")
    end_date = (date + timedelta(days=15)).strftime("%Y-%m-%d")

    search = catalog.search(
        collections=["sentinel-2-l2a"],
        bbox=bounds,
        datetime=f"{start_date}/{end_date}",
        query={"eo:cloud_cover": {"lt": 80}}  # filter cloudy scenes
    )

    items = list(search.items())
    if not items:
        return None

    signed_items = [planetary_computer.sign(item) for item in items]
    ds = odc.stac.load(
        signed_items,
        bands=["red", "green", "blue"],  # RGB bands
        bbox=bounds,
    )

    return ds.isel(time=0)

vanuatu_imagery = []

for idx, row in admin_boundaries_gdf.head(20).iterrows():
    polygon = row.geometry
    bbox = polygon.bounds
    event_date = '20/12/2024'
    formatted_date = datetime.strptime(event_date, "%d/%m/%Y").strftime("%Y-%m-%d")

    s1_data = fetch_s1_rtc_for_date_inference(bounds=list(bbox), date=formatted_date)
    s2_data = fetch_s2_rgb_for_date(bounds=list(bbox), date=formatted_date)

    if s1_data and s2_data:
        vanuatu_imagery.append([s1_data, s2_data])
        print("Added:", row.iname)

vanuatu_imagery

region_idx = 15 #10 
s1_data, s2_data = vanuatu_imagery[region_idx]

# Normalize VV
vv = s1_data["vv"].values
vv = (vv - vv.min()) / (vv.max() - vv.min())

# Normalize RGB
rgb = np.stack([s2_data["red"], s2_data["green"], s2_data["blue"]], axis=-1) #.values
rgb = (rgb - rgb.min()) / (rgb.max() - rgb.min())  # Normalize to 0–1

# Model prediction
img = torch.as_tensor(s1_data.to_array().values, device='cpu').unsqueeze(0)
output = model(img)
prediction = torch.softmax(output, dim=1).squeeze(0)
prediction = torch.argmax(prediction, dim=0).cpu().numpy()

# Resize prediction to match vv
if prediction.shape != vv.shape:
    prediction = resize(
        prediction, vv.shape, order=0, preserve_range=True, anti_aliasing=False
    ).astype(np.uint8)
    
# Define discrete colormap
cmap = mcolors.ListedColormap(["black", "green", "blue", "navy"])  # Colors for classes 0-3
bounds = [0, 1, 2, 3, 4]  # Boundaries between color levels
norm = mcolors.BoundaryNorm(bounds, cmap.N)

# Plot
fig, axes = plt.subplots(1, 3, figsize=(18, 6))

# Sentinel-2 RGB
axes[0].imshow(rgb)
axes[0].set_title("Sentinel-2 RGB")

# Sentinel-1 VV
axes[1].imshow(vv, cmap="gray")
axes[1].set_title("Sentinel-1 VV")

# Prediction with colormap
im = axes[2].imshow(prediction, cmap=cmap, norm=norm)
axes[2].set_title("Prediction")

# Adjust layout to make space for the legend
fig.subplots_adjust(bottom=0.2)

# Create a separate axis for the colorbar below the plots
cbar_ax = fig.add_axes([0.2, 0.08, 0.6, 0.03])  # [left, bottom, width, height]
cbar = fig.colorbar(im, cax=cbar_ax, orientation="horizontal")

# Fix the tick positions to be at class centers
tick_positions = [0.5, 1.5, 2.5, 3.5]  # Midpoints of the class ranges
cbar.set_ticks(tick_positions)
cbar.set_ticklabels([class_labels[i] for i in range(4)])

for ax in axes:
    ax.axis("off")

plt.show()

# Initialize a list to store results
vectorized_predictions = []

for idx, (s1_data, s2_data) in enumerate(vanuatu_imagery):
    polygon = admin_boundaries_gdf.iloc[idx].geometry  # Get the region geometry
    
    # Normalize Sentinel-1 VV
    vv = s1_data["vv"].values
    vv = (vv - vv.min()) / (vv.max() - vv.min())

    # Model Prediction
    img = torch.as_tensor(s1_data.to_array().values, device="cpu").unsqueeze(0)
    output = model(img)
    prediction = torch.softmax(output, dim=1).squeeze(0)
    prediction = torch.argmax(prediction, dim=0).cpu().numpy()

    # Vectorize Prediction using rasterio
    mask = prediction > 0  # Exclude "No Data" class (class 0)
    shapes_gen = shapes(prediction.astype(np.int16), mask=mask, transform=s1_data.rio.transform())

    # Convert raster shapes into GeoDataFrame features
    for geom, value in shapes_gen:
        class_id = int(value)
        vectorized_predictions.append({
            "geometry": shape(geom),
            "region": admin_boundaries_gdf.iloc[idx]["iname"],
            "class_id": class_id,
            "class_name": class_labels.get(class_id, "Unknown")
        })

# Create a GeoDataFrame
gdf_predictions = gpd.GeoDataFrame(vectorized_predictions) 
gdf_predictions.set_crs(epsg=32758, inplace=True) # WGS 84 / UTM zone 58S 

# Save as GeoJSON
gdf_predictions.to_file("vanuatu_flood_predictions.geojson", driver="GeoJSON")

print("GeoJSON file saved: vanuatu_flood_predictions.geojson")

s1_data.rio.crs

gdf_predictions.head(10)

# Define a colormap for class names
class_colors = {
    "No Data": "black",
    "Land (not flooded)": "green",
    "Flooded areas": "blue",
    "Permanent water bodies": "navy"
}

cmap = mcolors.ListedColormap(class_colors.values())
bounds = range(len(class_colors) + 1)
norm = mcolors.BoundaryNorm(bounds, cmap.N)

# Get unique regions
regions = gdf_predictions["region"].unique()
num_regions = len(regions)

# Set up grid layout (2 columns, rows calculated dynamically)
cols = 2
rows = math.ceil(num_regions / cols)  # Round up to fit all regions

# Increase figure size to make each subplot larger
fig, axes = plt.subplots(rows, cols, figsize=(16, rows * 8))  # Adjusted figure size

# Flatten axes array if there's more than 1 row
axes = axes.flatten() if num_regions > 1 else [axes]

for i, region in enumerate(regions):
    ax = axes[i]

    # Filter data for the specific region
    region_predictions = gdf_predictions[gdf_predictions["region"] == region]

    # Plot predictions
    region_predictions.plot(
        column="class_id",
        cmap=cmap,
        norm=norm,
        legend=False,  # Avoid duplicate legends
        ax=ax
    )

    ax.set_title(f"Flood Prediction for {region}", fontsize=18)
    ax.axis("off")

# Remove empty subplots if num_regions < rows * cols
for j in range(i + 1, len(axes)):
    fig.delaxes(axes[j])

# Add a single legend outside the plots
fig.legend(handles=ax.collections, labels=class_colors.keys(), loc="lower center", ncol=4, fontsize=14)

plt.tight_layout()
plt.show()