Flash flooding in rapidly growing urban centers like Nairobi poses a severe threat to infrastructure, informal settlements, and human life. Predicting these sudden inundation events requires blending static environmental vulnerability factors (elevation, slope, and land-cover imperviousness) with dynamic temporal triggers (real-time and historical precipitation patterns).
In this article, I will walk you through building a data-driven Flash Flood Early Warning System within Google Cloud's Colab Enterprise. We will use Google Earth Engine (GEE) to access satellite observations—specifically Global Precipitation Measurement (GPM) IMERG rainfall data and Shuttle Radar Topography Mission (SRTM) elevation models—train a PyTorch Long Short-Term Memory (LSTM) deep learning model to forecast flood triggers, and deploy the resulting model to Vertex AI Model Registry.
1. System Architecture & Flow
Before jumping into the notebook code, let's examine how Colab Enterprise bridges the gap between planetary-scale spatial data, deep learning accelerators, and production endpoints:
2. Environment Setup & Colab Enterprise Configuration
Colab Enterprise combines the familiar collaborative notebook interface with enterprise Google Cloud Platform (GCP) features: security governance, custom IAM service accounts, Vertex AI integration, and dedicated GPU runtimes.
Prerequisites & Permissions
GCP Project: An active Google Cloud Project with billing enabled.
IAM Roles: Ensure your user or service account has the following permissions:
- Earth Engine Resource Admin or Earth Engine Viewer
- Vertex AI User (roles/aiplatform.user)
- Storage Object Admin (for GCS artifact uploads)
Colab Enterprise Runtime: In your Colab Enterprise environment setup, select an NVIDIA T4 or L4 GPU runtime instance.
3. Step-by-Step Code Implementation
Environment Initialization & Ecosystem Authentication
This initial cell installs the required packages, initializes the Earth Engine API, and connects to the Vertex AI SDK.
# Install required libraries
!pip install -q geemap ee torch pandas scikit-learn google-cloud-aiplatform
import ee
import geemap
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from google.cloud import aiplatform
# 1. Authenticate and Initialize Google Earth Engine
try:
ee.Initialize()
except Exception as e:
ee.Authenticate()
ee.Initialize()
# 2. Initialize Vertex AI SDK
PROJECT_ID = "your-project-id" # <--- REPLACE WITH YOUR GCP PROJECT ID
REGION = "us-central1"
aiplatform.init(project=PROJECT_ID, location=REGION)
print("Environment Initialized. GEE and Vertex AI successfully connected.")
Spatial Area of Interest (AOI) & Static Data Extraction
Here we extract static environmental factors for Nairobi County: the Digital Elevation Model (SRTM DEM) to assess low-lying drainage basins, and the ESA WorldCover dataset to identify concrete/impervious surfaces that increase surface runoff.
# Define Nairobi County Bounding Box
nairobi_roi = ee.Geometry.Rectangle([36.6, -1.4, 37.1, -1.1])
def get_nairobi_static_data():
# 1. Digital Elevation Model (DEM) - SRTM 30m
dem = ee.Image("USGS/SRTMGL1_003").clip(nairobi_roi)
# 2. Land Use / Land Cover (LULC) - ESA WorldCover 10m
lulc = ee.ImageCollection("ESA/WorldCover/v100").first().clip(nairobi_roi)
# Sample pixels across Nairobi to construct spatial data frames
sample = dem.addBands(lulc).sample(region=nairobi_roi, scale=100, numPixels=1000)
df = geemap.ee_to_pandas(sample)
return df, dem, lulc
static_df, nairobi_dem, nairobi_lulc = get_nairobi_static_data()
print(f" Static Data Extracted. Sample size: {len(static_df)} spatial points.")
static_df.head()
Dynamic Time-Series Ingestion (GPM Precipitation)
Next, we query the Global Precipitation Measurement (GPM) IMERG dataset over a two-year window to extract historical rainfall time-series data for central Nairobi.
def get_rainfall_timeseries(start_date, end_date):
# GPM IMERG Final Precipitation L3 Half-Hourly
precipitation = (ee.ImageCollection("GPM/L3/IMERG_V06")
.filterBounds(nairobi_roi)
.filterDate(start_date, end_date)
.select('precipitationCal'))
# Extract time-series for Central Nairobi
nairobi_center = ee.Geometry.Point([36.8219, -1.2921])
ts = precipitation.getRegion(nairobi_center, 1000).getInfo()
# Format raw GEE JSON response into pandas DataFrame
df = pd.DataFrame(ts[1:], columns=[item[0] for item in ts[0]])
df = df[['time', 'precipitationCal']].sort_values('time')
df['time'] = pd.to_datetime(df['time'], unit='ms')
df['precipitationCal'] = pd.to_numeric(df['precipitationCal'], errors='coerce').fillna(0)
return df
# Fetch 2 years of rainfall history for model training
rain_df = get_rainfall_timeseries("2021-01-01", "2023-01-01")
print(f"Rainfall Time-series Extracted. Total records: {len(rain_df)}")
rain_df.head()
Sequence Windowing & Dataset
PreprocessingLSTMs require input formatted into sequence windows. We scale the precipitation values and construct rolling 72-hour historical windows. The target label $y=1$ represents an extreme precipitation event exceeding the 95th percentile.
def create_sequences(data, seq_length):
xs, ys = [], []
threshold = np.percentile(data, 95)
for i in range(len(data) - seq_length):
x = data[i:(i + seq_length)]
# Binary target: Trigger activated if next timestamp exceeds 95th percentile
y = 1 if data[i + seq_length] > threshold else 0
xs.append(x)
ys.append(y)
return np.array(xs), np.array(ys)
# 1. Scale precipitation values
scaler = StandardScaler()
rain_scaled = scaler.fit_transform(rain_df[['precipitationCal']]).flatten()
# 2. Build 72-period sliding windows
X, y = create_sequences(rain_scaled, seq_length=72)
# 3. Train/test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 4. Cast to PyTorch Tensors
X_train = torch.FloatTensor(X_train).unsqueeze(-1) # Shape: [batch, seq_len, features]
y_train = torch.FloatTensor(y_train).unsqueeze(-1)
X_test = torch.FloatTensor(X_test).unsqueeze(-1)
y_test = torch.FloatTensor(y_test).unsqueeze(-1)
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=32, shuffle=True)
print("Sequence Windowing Complete. Tensors prepared for GPU training.")
PyTorch LSTM Architecture & Training Loop
We define a 2-layer LSTM model that processes time-series precipitation windows and outputs a sigmoid flood trigger probability score.
class FloodLSTM(nn.Module):
def __init__(self):
super(FloodLSTM, self).__init__()
self.lstm = nn.LSTM(input_size=1, hidden_size=64, num_layers=2, batch_first=True)
self.fc = nn.Linear(64, 1)
self.sigmoid = nn.Sigmoid()
def forward(self, x):
out, _ = self.lstm(x)
out = self.fc(out[:, -1, :]) # Extract hidden state from last time step
return self.sigmoid(out)
# Instantiate model on GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = FloodLSTM().to(device)
criterion = nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# Model Training Loop
model.train()
for epoch in range(20):
for batch_x, batch_y in train_loader:
batch_x, batch_y = batch_x.to(device), batch_y.to(device)
optimizer.zero_grad()
outputs = model(batch_x)
loss = criterion(outputs, batch_y)
loss.backward()
optimizer.step()
if (epoch + 1) % 5 == 0:
print(f"Epoch {epoch+1}/20 | Loss: {loss.item():.4f}")
print("LSTM Model Training Finished.")
Spatial Inundation Mapping & Visualization
By coupling the LSTM's dynamic output probability with Nairobi's physical topography and land cover, we generate a high-resolution spatial inundation risk map directly within the notebook using geemap.
def generate_flood_risk_map(flood_probability):
# Calculate topography metrics
slope = ee.Terrain.slope(nairobi_dem)
# Normalize elevation (lower elevation values = higher vulnerability)
inv_elev = nairobi_dem.unitScale(0, 2500).multiply(-1).add(1)
# Mask impervious urban land surfaces (ESA WorldCover class 50 = Built-up)
urban_mask = nairobi_lulc.eq(50)
# Combined Spatial Risk Formula: Vulnerability Index * Dynamic Trigger Probability
risk_map = inv_elev.multiply(flood_probability).multiply(urban_mask)
return risk_map
# Simulated high-probability event forecast from LSTM (e.g., 85% probability)
predicted_event_probability = 0.85
nairobi_risk = generate_flood_risk_map(predicted_event_probability)
# Interactive Map Visualization inside Colab Enterprise
Map = geemap.Map()
Map.centerObject(nairobi_roi, 12)
Map.addLayer(
nairobi_risk,
{'min': 0, 'max': 1, 'palette': ['blue', 'yellow', 'red']},
'Flash Flood Spatial Risk Map'
)
Map
Registering the Model Artifact in Vertex AI
Finally, we export the trained PyTorch state dictionary and upload it into the Vertex AI Model Registry for future serving and enterprise management.
import os
# 1. Save trained PyTorch model state locally
model.cpu()
os.makedirs("model_artifacts", exist_ok=True)
torch.save(model.state_dict(), "model_artifacts/nairobi_flood_lstm.pth")
# 2. Upload artifact to Vertex AI Model Registry
BUCKET_URI = f"gs://{PROJECT_ID}-flood-models" # Ensure GCS bucket exists
# Copy artifact to Cloud Storage
!gsutil cp model_artifacts/nairobi_flood_lstm.pth {BUCKET_URI}/flood-model/
model_registry = aiplatform.Model.upload(
display_name="Nairobi-Flash-Flood-LSTM",
artifact_uri=f"{BUCKET_URI}/flood-model/",
serving_container_image_uri="us-docker.pkg.dev/vertex-ai/prediction/pytorch-cpu.1-13:latest",
)
print(f"Model successfully registered to Vertex AI: {model_registry.resource_name}")
Major Architectural Takeaways for this project
Decoupled Data and Compute: By leveraging Google Earth Engine directly inside Colab Enterprise, heavy geospatial raster algebra runs on Google's cloud infrastructure without overwhelming local notebook memory.
Hybrid Modeling Strategy: Combining temporal deep learning models (LSTMs) for trigger prediction with physical spatial rasters (SRTM elevation and ESA land cover) yields actionable, highly targeted risk maps.
Path to Enterprise Production: Storing weights in Google Cloud Storage and registering models in Vertex AI provides a clear transition path from research sandbox to operational API endpoints.
Top comments (0)