Added Surrogate Modeling enhancement module (#30)

* update the benchmark fim_lookup query within fimserv

* update the date field and  formatting of codes

* added the evaluation module within FIMserv

* added the surrogate modeling enhacement module

* added Surrogate Model citation

Author: supathdhitalGEO

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "fimserve"
-version = "0.1.87"
+version = "0.1.88"
 description = "Framework which is developed with the purpose of quickly generating Flood Inundation Maps (FIM) for emergency response and risk assessment. It is developed under Surface Dynamics Modeling Lab (SDML)."
 authors = [{ name = "Supath Dhital", email = "sdhital@crimson.ua.edu" }]
 maintainers = [{ name = "Supath Dhital", email = "sdhital@crimson.ua.edu" }]

src/fimserve/enhancement_withSM/README.md

@@ -0,0 +1,44 @@
+# Operational Enhanced FIM
+
+## Overview
+
+The **Operational Enhanced Flood Inundation Mapping (FIM)** framework provides a streamlined pipeline for real-time flood mapping and impact analysis. Its primary objective is to support emergency response operations by delivering timely, high-resolution flood information.
+
+
+## Workflow
+
+The core components of the framework are illustrated in the figure below:
+
+![Workflow](graphics/operationalFIM.jpg)
+
+## Key Capabilities
+
+- **Real-time flood map generation** using surrogate models.
+- **Automated processing pipeline** for rapid data acquisition, model execution, and visualization.
+- **Scalable architecture** adaptable to different regions and data sources.
+
+## Exposure Analysis
+
+In addition to mapping flooded areas, the framework includes modules for:
+
+- **Population exposure estimation**: Identifies the number of people affected by the flood event. For this it uses the population grids to analyze the exposed population. The framework provides the estimated exposed population and the spatial distribution of those population count using a unique way of vizualization.
+![Workflow](graphics/population_exposure.png)
+
+- **Building exposure analysis**: Detects flooded structures using geospatial building footprints and flood maps.
+![Workflow](graphics/building_exposure.png)
+
+These capabilities help quantify potential impacts and support decision-making in emergency situations.
+
+## Application in Emergency Response
+
+This framework is designed to:
+- Provide first responders with up-to-date flood information
+- Inform public warnings and evacuation strategies
+- Aid in post-event damage assessments and recovery planning
+
+## Reference
+The trained Surrogate Model that is being applied for the enhancement of the FIMserv derieved Low-Fidelity results comes out of the following research, which is currently in review. 
+
+**Preprint** — **Supath Dhital**, Sagy Cohen, Parvaneh Nikrou, et al.  
+Enhancement of low-fidelity flood inundation mapping through surrogate modeling. *ESS Open Archive*, November 03, 2025.  
+[https://doi.org/10.22541/essoar.176218121.12875584/v1](https://doi.org/10.22541/essoar.176218121.12875584/v1)

src/fimserve/enhancement_withSM/SM_prediction.py

@@ -0,0 +1,170 @@
+from SM_preprocess import *
+from surrogate_model import *
+from utlis import *
+from preprocessFIM import *
+
+def load_model(model):
+    # Set up S3 access
+    fs = s3fs.S3FileSystem(anon=True)
+    bucket_path = "sdmlab/SM_dataset/trained_model/SM_trainedmodel.ckpt"
+
+    # Download to a temporary file
+    with fs.open(bucket_path, 'rb') as s3file:
+        with tempfile.NamedTemporaryFile(suffix=".ckpt", delete=False) as tmp_ckpt:
+            tmp_ckpt.write(s3file.read())
+            tmp_ckpt_path = tmp_ckpt.name
+
+    # Load checkpoint
+    checkpoint = torch.load(tmp_ckpt_path, map_location='cuda' if torch.cuda.is_available() else 'cpu')
+    model.load_state_dict(checkpoint['state_dict'])
+
+    # Move model to device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model.to(device)
+    model.eval()
+
+    return model, device
+
+#WEIGHTED AVERAGE for patch
+def create_weight_map(M: int, N: int):
+    weight_map = np.zeros((M, N), dtype=np.float32)
+    center_x, center_y = M // 2, N // 2
+    for i in range(M):
+        for j in range(N):
+            dist_sq = (i - center_x)**2 + (j - center_y)**2
+            weight = np.exp(-dist_sq / (2 * (min(M, N) / 2)**2))
+            weight_map[i, j] = weight
+    return torch.from_numpy(weight_map).float().unsqueeze(0).unsqueeze(0) # (1, 1, M, N)
+
+#If there is Stride
+def predict_on_area(dataset, model, shape: torch.Tensor, M: int = 256, N: int = 256, stride: int = 128, device=None):
+    # Get row and col size
+    shape_row = shape.size(1)
+    shape_col = shape.size(2)
+
+    # Pad if needed
+    pad_h = (stride * ((shape_row - M) // stride + 1) + M - shape_row)
+    pad_w = (stride * ((shape_col - N) // stride + 1) + N - shape_col)
+
+    if pad_h > 0 or pad_w > 0:
+        padding = (0, pad_w, 0, pad_h) 
+        shape = nn.functional.pad(shape, padding, mode='constant', value=0)
+
+    # Update new shape after padding
+    new_row = shape.size(1)
+    new_col = shape.size(2)
+
+    # Separate X and Y
+    X = shape[dataset.x_feature_index]
+    y = shape[dataset.y_feature_index]
+
+    # Initialize weighted prediction sum and weight sum arrays
+    weighted_prediction_sum = torch.zeros((1, new_row, new_col), device=device)
+    weight_sum = torch.zeros((1, new_row, new_col), device=device)
+
+    # Create the weight map
+    weight_map = create_weight_map(M, N).to(device)
+
+    # Loop over patches
+    for start_i in range(0, new_row - M + 1, stride):
+        for start_j in range(0, new_col - N + 1, stride):
+            end_i = start_i + M
+            end_j = start_j + N
+            patch = X[:, start_i:end_i, start_j:end_j].unsqueeze(0).to(device)
+
+            with torch.no_grad():
+                patch_prediction_raw = model(patch)
+
+            weighted_prediction = patch_prediction_raw * weight_map
+            weighted_prediction_sum[:, start_i:end_i, start_j:end_j] += weighted_prediction.squeeze(0) 
+            weight_sum[:, start_i:end_i, start_j:end_j] += weight_map.squeeze(0)
+
+    epsilon = 1e-8
+    final_prediction = weighted_prediction_sum / (weight_sum + epsilon)
+    final_prediction = (final_prediction > 0.01).float() 
+
+    # Crop back to original shape (before padding)
+    final_prediction = final_prediction[:, :shape_row, :shape_col]
+    y = y[:, :shape_row, :shape_col]
+    lf = shape[[dataset.lf_index]][:, :shape_row, :shape_col]
+
+    return final_prediction.cpu(), y.cpu(), lf.cpu()
+
+#Save the tif file 
+def save_image(image: torch.Tensor, path: Path, reference_tif: str):
+    """Save the image as a .tif file.
+    
+    Args:
+        image (torch.Tensor): The image to save
+        path (Path): The path to save the image
+    """
+    image_np = image.squeeze().cpu().numpy().astype('float32')
+    with rasterio.open(reference_tif) as ref:
+        meta = ref.meta.copy()
+        meta.update({
+            "driver": "GTiff",
+            "height": image_np.shape[0],
+            "width": image_np.shape[1],
+            "count": 1,
+            "dtype": 'float32'
+        })
+
+        with rasterio.open(path, 'w', **meta) as dst:
+            dst.write(image_np, 1)
+        mask_with_PWB(path, path)
+
+        with rasterio.open(path, 'r+') as dst:
+            data = dst.read(1)
+            binary_data = np.where(data > 0, 1, 0).astype(np.uint8)
+            dst.write(binary_data, 1)
+
+        compress_tif_lzw(path)
+        
+        
+#ENHANCE THE LOW-FIDELITY FLOOD MAP
+def Predict_FM(huc_id, patch_size=(256, 256)):
+    
+    data_dir = Path(f'./HUC{huc_id}_forcings/')
+    model =  AttentionUNet(channel=8)
+    
+    preprocessor = InferenceDataPreprocessor(data_dir=Path(data_dir), patch_size=patch_size, verbose=True)
+    
+    print("Loading model...")
+    model, device = load_model(model)
+    print("Model loaded.")
+    
+    
+    lf_files = preprocessor.get_all_lf_maps(huc_id)
+    for lf_path in lf_files:
+        lf_filename = lf_path.name
+        print(f"Predicting for: {lf_filename}\n")
+
+        print(f"Loading static features for HUC {huc_id}...")
+        static_stack = preprocessor.get_static_stack(huc_id)
+        lf_tensor = preprocessor.tif_to_tensor(lf_path, feature_name='low_fidelity')
+
+        # Combine and validate
+        area_tensor = torch.cat([static_stack, lf_tensor], dim=0)
+        if area_tensor.shape[0] != 8:
+            raise ValueError(f"Expected 8 channels, got {area_tensor.shape[0]} — check missing static feature for HUC {huc_id}.")
+
+        # Define dummy interface
+        class Dummy:
+            x_feature_index = list(range(area_tensor.shape[0])) 
+            y_feature_index = [area_tensor.shape[0] - 1]     
+            lf_index = area_tensor.shape[0] - 1
+
+        print(f"Static features loaded for {huc_id}.\n")
+        
+        # Predict
+        print(f"Enhancing {lf_path}...")
+        x, y, lf = predict_on_area(Dummy, model, area_tensor, M=patch_size[0], N=patch_size[1], stride=patch_size[0] // 2, device=device)
+
+        # Save result
+        pred_dir = Path(f"./Results/HUC{huc_id}/")
+        pred_dir.mkdir(parents=True, exist_ok=True)
+        pred_path = pred_dir / f"SMprediction_{lf_filename}"
+        save_image(x, pred_path, lf_path)
+        print(f"Enhancement completed for {lf_filename}.\n")
+        
+

src/fimserve/enhancement_withSM/SM_preprocess.py

@@ -0,0 +1,123 @@
+import torch
+import rasterio
+import numpy as np
+from pathlib import Path
+import torch.nn as nn
+from scipy.stats import boxcox
+
+GLOBAL_STATS = {
+    'elevation': {'mean': 651.62, 'std': 935.30},
+    'slope': {'mean': 1.65, 'std': 2.50},
+    'flow_acc': {'mean': -101.88, 'std': 65.45},
+    'twi': {'mean': -4.23, 'std': 2.27, 'min': -11.01, 'max': 25.39},
+    'curve_number': {'mean': 73.96, 'std': 13.29, 'min': 30.12, 'max': 100.00},
+    'soil_moisture': {'mean': 28.06, 'std': 9.94, 'min': 0.00, 'max': 76.13},
+}
+
+# INFERENCE DATA PREPROCESSOR
+class InferenceDataPreprocessor:
+    
+    STATIC_FEATURES = ['curve_number', 'elevation', 'flow_acc', 'lulc', 'slope', 'soil_moisture', 'twi']
+    LF_KEYWORDS = ['hand']
+
+    FEATURE_FILENAME_MAP = {
+        'curve_number': 'CN',
+        'flow_acc': 'flowacc',
+        'soil_moisture': 'SM',
+        'elevation': 'elevation',
+        'slope': 'slope',
+        'twi': 'twi',
+        'lulc': 'LULC'
+    }
+
+    def __init__(self, data_dir: Path, patch_size=(128, 128), global_stats=None, verbose=False):
+        self.data_dir = Path(data_dir)
+        self.M, self.N = patch_size
+        self.verbose = verbose
+        self.global_stats = global_stats if global_stats else GLOBAL_STATS
+
+    def tif_to_tensor(self, path: Path, feature_name: str = None) -> torch.Tensor:
+        with rasterio.open(path) as src:
+            array = src.read(1).astype(np.float32)
+            nodata_value = src.nodata
+            if nodata_value is not None:
+                array[array == nodata_value] = np.nan
+            array = np.nan_to_num(array, nan=0.0)
+            tensor = torch.tensor(array, dtype=torch.float32).unsqueeze(0)
+
+            if feature_name == 'elevation':
+                mean, std = self.global_stats['elevation'].values()
+                tensor = (tensor - mean) / (std + 1e-7)
+            elif feature_name == 'slope':
+                tensor = self.apply_boxcox(tensor)
+                mean, std = self.global_stats['slope'].values()
+                tensor = (tensor - mean) / (std + 1e-7)
+            elif feature_name == 'flow_acc':
+                tensor = self.apply_boxcox(tensor)
+                mean, std = self.global_stats['flow_acc'].values()
+                tensor = (tensor - mean) / (std + 1e-7)
+            elif feature_name == 'lulc':
+                reclass_map = {1: 1, 2: 2, 4: 3, 3: 4, 8: 5, 6: 6, 7: 7, 5: 8, 9: 9}
+                array = array.astype(np.int32)
+                reclass_array = np.vectorize(lambda x: reclass_map.get(x, 0))(array).astype(np.float32)
+                tensor = torch.tensor(reclass_array, dtype=torch.float32).unsqueeze(0)
+            elif feature_name == 'low_fidelity':
+                tensor = (tensor > 0).float()
+            elif feature_name in self.global_stats:
+                min_val = self.global_stats[feature_name]['min']
+                max_val = self.global_stats[feature_name]['max']
+                tensor = (tensor - min_val) / (max_val - min_val + 1e-7)
+                
+            return tensor
+
+    def apply_boxcox(self, tensor: torch.Tensor, lmbda=0.5) -> torch.Tensor:
+        tensor = tensor + 1e-6
+        flat_np = tensor.flatten().numpy()
+        transformed = boxcox(flat_np, lmbda=lmbda)
+        return torch.tensor(transformed).reshape(tensor.shape)
+
+    def patchify(self, data: torch.Tensor):
+        C, H, W = data.shape
+        stride_h = stride_w = self.M // 2
+        pad_h = (stride_h * ((H - self.M) // stride_h + 1) + self.M - H) % stride_h
+        pad_w = (stride_w * ((W - self.N) // stride_w + 1) + self.N - W) % stride_w
+        padded = nn.functional.pad(data, (0, int(pad_w), 0, int(pad_h)), mode='constant', value=0)
+        patches = padded.unfold(1, self.M, stride_h).unfold(2, self.N, stride_w)
+        patches = patches.permute(1, 2, 0, 3, 4).reshape(-1, C, self.M, self.N)
+        return patches
+
+    def get_static_stack(self, huc_id: str):
+        tensors = []
+        for feature in self.STATIC_FEATURES:
+            search_key = self.FEATURE_FILENAME_MAP[feature]
+            match = list(self.data_dir.glob(f"*{search_key}*{huc_id}*.tif"))
+            if not match:
+                print(f"Missing static feature: {feature} for {huc_id}")
+                return None
+            tensors.append(self.tif_to_tensor(match[0], feature_name=feature))
+        return torch.cat(tensors, dim=0)
+
+    def get_all_lf_maps(self, huc_id: str):
+        return sorted([
+            f for f in self.data_dir.glob(f"*{huc_id}*.tif")
+            if any(k in f.name.lower() for k in self.LF_KEYWORDS)
+        ])
+
+    def preprocess_all_lf_maps(self, huc_id: str):
+        static_stack = self.get_static_stack(huc_id)
+        if static_stack is None:
+            return []
+
+        lf_files = self.get_all_lf_maps(huc_id)
+        results = []
+
+        for lf_path in lf_files:
+            lf_tensor = self.tif_to_tensor(lf_path, feature_name='low_fidelity')
+            combined = torch.cat([static_stack, lf_tensor], dim=0)
+            patches = self.patchify(combined)
+            results.append((lf_path.name, patches, lf_path))
+            if self.verbose:
+                print(f"Processed {lf_path.name} with {patches.shape[0]} patches.")
+
+        return results
+