Upload: load.py

single/load.py

@@ -1,15 +1,32 @@
+"""
+load.py
+
+Module for loading ensemble models (STAC compatible) and performing
+optimized inference on large geospatial imagery using dynamic batching
+and Gaussian blending.
+"""
+
+import math
+import pathlib
+import itertools
+from typing import Literal, Tuple, List
+
 import torch
 import torch.nn
-import pathlib
-import pystac
-from typing import Literal, Tuple
 import numpy as np
-import itertools
+import pystac
+from torch.utils.data import Dataset, DataLoader
 from tqdm import tqdm
-import math
 
-# Ensemble model for combining multiple models' outputs
+# ==============================================================================
+# 1. HELPER CLASSES & FUNCTIONS
+# ==============================================================================
+
 class EnsembleModel(torch.nn.Module):
+    """
+    Runtime ensemble model for combining multiple model outputs.
+    Used when loading multiple separate .pt2 files.
+    """
     def __init__(self, *models, mode="max"):
         super(EnsembleModel, self).__init__()
         self.models = torch.nn.ModuleList(models)
@@ -19,58 +36,114 @@ class EnsembleModel(torch.nn.Module):
 
     def forward(self, x) -> Tuple[torch.Tensor, torch.Tensor]:
         """
-        Forward pass for ensemble.
-        
         Returns:
-            Tuple of (probabilities, uncertainty):
-            - probabilities: (B, 1, H, W) - aggregated predictions
-            - uncertainty: (B, 1, H, W) - normalized std deviation
+            - probabilities: (B, 1, H, W)
+            - uncertainty: (B, 1, H, W) (normalized std dev)
         """
-        outputs = []
-        for model in self.models:
-            output = model(x)
-            outputs.append(output)
+        outputs = [model(x) for model in self.models]
         
         if not outputs:
             return None, None
         
-        # Stack all model outputs: (B, N, H, W) where N = number of models
-        stacked_outputs = torch.stack(outputs, dim=1)  # (B, N, 1, H, W)
-        stacked_outputs = stacked_outputs.squeeze(2)   # (B, N, H, W)
+        # Stack: (B, N, H, W)
+        stacked = torch.stack(outputs, dim=1).squeeze(2)
         
-        # Calculate aggregated probabilities
+        # Aggregation
         if self.mode == "max":
-            output_probs = torch.max(stacked_outputs, dim=1, keepdim=True)[0]
+            probs = torch.max(stacked, dim=1, keepdim=True)[0]
         elif self.mode == "mean":
-            output_probs = torch.mean(stacked_outputs, dim=1, keepdim=True)
+            probs = torch.mean(stacked, dim=1, keepdim=True)
         elif self.mode == "median":
-            output_probs = torch.median(stacked_outputs, dim=1, keepdim=True)[0]
+            probs = torch.median(stacked, dim=1, keepdim=True)[0]
         elif self.mode == "min":
-            output_probs = torch.min(stacked_outputs, dim=1, keepdim=True)[0]
+            probs = torch.min(stacked, dim=1, keepdim=True)[0]
         elif self.mode == "none":
-            # Return all predictions without aggregation
-            return stacked_outputs, None
-        else:
-            raise ValueError("Mode must be 'min', 'mean', 'median', or 'max'.")
-        
-        # Calculate uncertainty (normalized standard deviation)
+            return stacked, None
+            
+        # Uncertainty
         N = len(outputs)
         if N > 1:
-            # Calculate std across models (dim=1)
-            std_output = torch.std(stacked_outputs, dim=1, keepdim=True)
-            
-            # Normalize the standard deviation [0 - 1]
-            # Formula: std_max = sqrt(0.25 * N / (N - 1))
+            std = torch.std(stacked, dim=1, keepdim=True)
             std_max = math.sqrt(0.25 * N / (N - 1))
-            uncertainty = std_output / std_max
-            
-            # Clamp to [0, 1] to avoid numerical issues
-            uncertainty = torch.clamp(uncertainty, 0.0, 1.0)
+            uncertainty = torch.clamp(std / std_max, 0.0, 1.0)
         else:
-            # Single model: no uncertainty
-            uncertainty = torch.zeros_like(output_probs)
+            uncertainty = torch.zeros_like(probs)
         
-        return output_probs, uncertainty  # Both (B, 1, H, W)
+        return probs, uncertainty
+
+def get_spline_window(window_size: int, power: int = 2) -> np.ndarray:
+    """Generates a 2D Hann window for smoothing tile edges."""
+    intersection = np.hanning(window_size)
+    window_2d = np.outer(intersection, intersection)
+    return (window_2d ** power).astype(np.float32)
+
+def fix_lastchunk(iterchunks, s2dim, chunk_size):
+    """Adjusts the last chunks to fit within image boundaries."""
+    itercontainer = []
+    for index_i, index_j in iterchunks:
+        if index_i + chunk_size > s2dim[0]:
+            index_i = max(s2dim[0] - chunk_size, 0)
+        if index_j + chunk_size > s2dim[1]:
+            index_j = max(s2dim[1] - chunk_size, 0)
+        itercontainer.append((index_i, index_j))
+    return list(set(itercontainer))
+
+def define_iteration(dimension: tuple, chunk_size: int, overlap: int = 0):
+    """Generates top-left coordinates for sliding window inference."""
+    dimy, dimx = dimension
+    if chunk_size > max(dimx, dimy):
+        return [(0, 0)]
+        
+    y_step = chunk_size - overlap
+    x_step = chunk_size - overlap
+    
+    iterchunks = list(itertools.product(
+        range(0, dimy, y_step), 
+        range(0, dimx, x_step)
+    ))
+    
+    return fix_lastchunk(iterchunks, dimension, chunk_size)
+
+# ==============================================================================
+# 2. DATASET FOR PARALLEL LOADING
+# ==============================================================================
+
+class PatchDataset(Dataset):
+    """
+    Dataset wrapper to handle image slicing and padding on CPU workers.
+    """
+    def __init__(self, image: np.ndarray, coords: List[Tuple[int, int]], chunk_size: int, nodata: float = 0):
+        self.image = image
+        self.coords = coords
+        self.chunk_size = chunk_size
+        self.nodata = nodata
+
+    def __len__(self):
+        return len(self.coords)
+
+    def __getitem__(self, idx):
+        row_off, col_off = self.coords[idx]
+        
+        # Read patch
+        patch = self.image[:, row_off : row_off + self.chunk_size, col_off : col_off + self.chunk_size]
+        c, h, w = patch.shape
+
+        patch_tensor = torch.from_numpy(patch).float()
+
+        # Apply padding if patch is smaller than chunk_size (edges)
+        pad_h = self.chunk_size - h
+        pad_w = self.chunk_size - w
+        if pad_h > 0 or pad_w > 0:
+            patch_tensor = torch.nn.functional.pad(patch_tensor, (0, pad_w, 0, pad_h), "constant", self.nodata)
+
+        # Identify nodata pixels
+        mask_nodata = (patch_tensor == self.nodata).all(dim=0)
+        
+        return patch_tensor, row_off, col_off, h, w, mask_nodata
+
+# ==============================================================================
+# 3. LOADING & INFERENCE LOGIC
+# ==============================================================================
 
 def compiled_model(
     path: pathlib.Path, 
@@ -79,238 +152,160 @@ def compiled_model(
     *args, **kwargs
 ):
     """
-    Loads model(s) dynamically based on STAC metadata.
-    
-    - If single .pt2 → returns single model
-    - If multiple .pt2 → returns EnsembleModel
-    
-    Args:
-        mode: Aggregation mode for ensembles (ignored for single models)
-    
-    Returns:
-        Single model or EnsembleModel
+    Loads .pt2 model(s). Returns a single model or an EnsembleModel.
+    Automatically unwraps ExportedProgram if possible.
     """
-    model_paths = []
-    for asset_key, asset in stac_item.assets.items():
-        if asset.href.endswith(".pt2"):
-            model_paths.append(asset.href) 
+    model_paths = sorted([
+        asset.href for key, asset in stac_item.assets.items() 
+        if asset.href.endswith(".pt2")
+    ])
     
     if not model_paths:
         raise ValueError("No .pt2 files found in STAC item assets.")
     
-    model_paths.sort()
-    
+    # Helper to load and unwrap
+    def load_pt2(p):
+        program = torch.export.load(p)
+        return program.module() if hasattr(program, "module") else program
+
     if len(model_paths) == 1:
-        # Single model
-        return torch.export.load(model_paths[0]).module()
+        return load_pt2(model_paths[0])
     else:
-        # Ensemble model
-        models = [torch.export.load(p).module() for p in model_paths]
+        models = [load_pt2(p) for p in model_paths]
         return EnsembleModel(*models, mode=mode)
 
-def define_iteration(dimension: tuple, chunk_size: int, overlap: int = 0):
-    """
-    Defines iteration strategy to traverse the image with overlap.
-    """
-    dimy, dimx = dimension
-    if chunk_size > max(dimx, dimy):
-        return [(0, 0)]
-    y_step = chunk_size - overlap
-    x_step = chunk_size - overlap
-    iterchunks = list(itertools.product(range(0, dimy, y_step), range(0, dimx, x_step)))
-    iterchunks_fixed = fix_lastchunk(
-        iterchunks=iterchunks, s2dim=dimension, chunk_size=chunk_size
-    )
-    return iterchunks_fixed
-
-
-def fix_lastchunk(iterchunks, s2dim, chunk_size):
-    """
-    Adjusts last chunks to prevent them from exceeding boundaries.
-    """
-    itercontainer = []
-    for index_i, index_j in iterchunks:
-        if index_i + chunk_size > s2dim[0]:
-            index_i = max(s2dim[0] - chunk_size, 0)
-        if index_j + chunk_size > s2dim[1]:
-            index_j = max(s2dim[1] - chunk_size, 0)
-        itercontainer.append((index_i, index_j))
-    return list(set(itercontainer)) # Returns unique values just in case
-
 
 def predict_large(
     image: np.ndarray,
     model: torch.nn.Module,
     chunk_size: int = 512,
-    overlap: int = 64,
-    device: str = "cpu",
+    overlap: int = 128,     
+    batch_size: int = 16, 
+    num_workers: int = 8,    # Recommended: 8-16
+    device: str = "cuda",
     nodata: float = 0.0
 ) -> Tuple[np.ndarray, np.ndarray] | np.ndarray:
     """
-    Predict a full 'image' (C, H, W) using overlapping patches.
-    
-    Args:
-        image: Input array (C, H, W)
-        model: Compiled PyTorch model
-        chunk_size: Tile size for inference
-        overlap: Overlap between tiles
-        device: 'cpu' or 'cuda'
-        nodata: No-data value
-        
-    Returns:
-        - For ensembles: Tuple of (probabilities, uncertainty), both (1, H, W)
-        - For single models: probabilities array (1, H, W)
-    
-    Compatible with:
-    - Normal models (with .eval()) - returns probabilities only
-    - Exported models (.pt2) - returns probabilities only
-    - Ensembles (EnsembleModel) - returns (probabilities, uncertainty)
+    Optimized inference for large images using Dynamic Batching and Gaussian Blending.
     """
     
-    # Validate input array dimensions
     if image.ndim != 3:
-        raise ValueError(f"Input array must be (C, H, W). Received {image.shape}")
+        raise ValueError(f"Input image must be (C, H, W). Received {image.shape}")
     
-    bands, height, width = image.shape
-    
-    # Prepare model (compatibility logic for .pt2 models)
+    # --- 1. Robust Model Unwrapping ---
+    # Fix for torch.export.load() returning an ExportedProgram container
+    if hasattr(model, "module") and callable(model.module):
+        try:
+            unpacked = model.module()
+            if isinstance(unpacked, torch.nn.Module):
+                model = unpacked
+        except Exception:
+            pass
+
+    # --- 2. Setup Model ---
     try:
         model.eval()
-        for p in model.parameters():
-            p.requires_grad = False
-        model = model.to(device)
-    except (NotImplementedError, AttributeError):
-        # Exported model (.pt2) or EnsembleModel
+        for p in model.parameters(): p.requires_grad = False
+    except: pass
+    
+    # Only move to device if it's a standard Module (ExportedProgram handles device internally or via input)
+    if isinstance(model, torch.nn.Module):
         model = model.to(device)
-        
-    test_input = torch.zeros(1, bands, chunk_size, chunk_size).to(device)
+
+    bands, height, width = image.shape
+
+    # --- 3. Check Signature (Ensemble vs Single) ---
+    # Dummy pass (batch=2 to respect dynamic shapes)
+    dummy = torch.randn(2, bands, chunk_size, chunk_size).to(device)
     with torch.no_grad():
-        test_output = model(test_input)
-    
-    is_ensemble = isinstance(test_output, tuple) and len(test_output) == 2
-    
-    # Initialize output arrays
-    output_probs = np.full((1, height, width), nodata, dtype=np.float32)
-    
-    if is_ensemble:
-        output_uncertainty = np.full((1, height, width), nodata, dtype=np.float32)
+        out = model(dummy)
+    is_ensemble = isinstance(out, tuple) and len(out) == 2
+
+    # --- 4. Initialize Buffers (Accumulators) ---
+    out_probs = np.zeros((1, height, width), dtype=np.float32)
+    count_map = np.zeros((1, height, width), dtype=np.float32)
+    out_uncert = np.zeros((1, height, width), dtype=np.float32) if is_ensemble else None
+
+    # --- 5. Prepare Spline Window ---
+    window_spline = get_spline_window(chunk_size, power=2)
+    window_tensor = torch.from_numpy(window_spline).to(device)
 
-    # Get the list of tile offsets
-    coords = define_iteration(
-        dimension=(height, width),
-        chunk_size=chunk_size,
-        overlap=overlap
+    # --- 6. DataLoader Setup ---
+    coords = define_iteration((height, width), chunk_size, overlap)
+    dataset = PatchDataset(image, coords, chunk_size, nodata)
+    loader = DataLoader(
+        dataset, 
+        batch_size=batch_size, 
+        shuffle=False, 
+        num_workers=num_workers,
+        prefetch_factor=2, 
+        pin_memory=True
     )
 
-    # Iterate over tiles
-    for idx, (row_off, col_off) in enumerate(tqdm(coords, desc="Inference")):
-        
-        # Read chunk (numpy slicing)
-        patch = image[
-            :,
-            row_off : row_off + chunk_size,
-            col_off : col_off + chunk_size
-        ]
-
-        # Convert to tensor and handle padding if tile is smaller than chunk_size
-        patch_tensor = torch.from_numpy(patch).float().unsqueeze(0).to(device)
-        _, _, h_tile, w_tile = patch_tensor.shape
-        
-        # Calculate padding needed
-        pad_h = chunk_size - h_tile
-        pad_w = chunk_size - w_tile
-        
-        # Apply padding if necessary
-        if pad_h > 0 or pad_w > 0:
-            patch_tensor = torch.nn.functional.pad(
-                patch_tensor, (0, pad_w, 0, pad_h), "constant", nodata
-            )
+    # --- 7. Inference Loop ---
+    for batch in tqdm(loader, desc=f"Inference (Batch {batch_size})"):
+        patches, r_offs, c_offs, h_actuals, w_actuals, nodata_masks = batch
         
-        # Create mask for nodata areas (all bands are nodata)
-        mask_all = (patch_tensor == nodata).all(dim=1, keepdim=True)
-        
-        # Forward pass
+        # Move inputs to GPU
+        patches = patches.to(device, non_blocking=True)
+        nodata_masks = nodata_masks.to(device, non_blocking=True) # (B, H, W)
+
+        # Forward Pass
         with torch.no_grad():
-            model_output = model(patch_tensor)
-            
             if is_ensemble:
-                probs, uncertainty = model_output
-                probs = probs.masked_fill(mask_all, nodata)
-                uncertainty = uncertainty.masked_fill(mask_all, nodata)
+                probs, uncert = model(patches)
             else:
-                probs = model_output
-                probs = probs.masked_fill(mask_all, nodata)
-        
-        # Remove batch dimension and ensure (1, H, W)
-        if probs.ndim == 4:
-            probs = probs.squeeze(0)  # (1, H, W)
-        
-        # Convert to numpy
-        result_probs = probs.cpu().numpy()  # (1, H, W)
-        
-        if is_ensemble:
-            if uncertainty.ndim == 4:
-                uncertainty = uncertainty.squeeze(0)
-            result_uncertainty = uncertainty.cpu().numpy()
+                probs = model(patches)
+                uncert = None
 
-        # Logic for partial writing
-        if col_off == 0:
-            offset_x = 0
-        else:
-            offset_x = col_off + overlap // 2
+        # Ensure correct dimensions (B, C, H, W)
+        if probs.ndim == 3: probs = probs.unsqueeze(1)
+        if is_ensemble and uncert.ndim == 3: uncert = uncert.unsqueeze(1)
 
-        if row_off == 0:
-            offset_y = 0
-        else:
-            offset_y = row_off + overlap // 2
+        # Prepare weights for batch
+        B = patches.size(0)
+        batch_weights = window_tensor.unsqueeze(0).unsqueeze(0).repeat(B, 1, 1, 1)
+        
+        # Zero out weights where input was nodata
+        batch_weights[nodata_masks.unsqueeze(1)] = 0.0
 
-        if (offset_x + chunk_size) == width:
-            length_x = chunk_size
-            sub_x_start = 0
-        else:
-            length_x = chunk_size - (overlap // 2)
-            sub_x_start = overlap // 2 if col_off != 0 else 0
+        # Apply weights
+        probs_weighted = probs * batch_weights
+        if is_ensemble:
+            uncert_weighted = uncert * batch_weights
 
-        if (offset_y + chunk_size) == height:
-            length_y = chunk_size
-            sub_y_start = 0
-        else:
-            length_y = chunk_size - (overlap // 2)
-            sub_y_start = overlap // 2 if row_off != 0 else 0
-
-        # Ensure we don't exceed array bounds
-        if offset_y + length_y > height:
-            length_y = height - offset_y
-        if offset_x + length_x > width:
-            length_x = width - offset_x
-
-        # Extract the valid region from the result
-        to_write_probs = result_probs[
-            :,
-            sub_y_start : sub_y_start + length_y,
-            sub_x_start : sub_x_start + length_x
-        ]
-
-        # Write to the output numpy array
-        output_probs[
-            :,
-            offset_y : offset_y + length_y,
-            offset_x : offset_x + length_x
-        ] = to_write_probs
-        
+        # Move to CPU
+        probs_cpu = probs_weighted.cpu().numpy()
+        weights_cpu = batch_weights.cpu().numpy()
         if is_ensemble:
-            to_write_uncertainty = result_uncertainty[
-                :,
-                sub_y_start : sub_y_start + length_y,
-                sub_x_start : sub_x_start + length_x
-            ]
-            output_uncertainty[
-                :,
-                offset_y : offset_y + length_y,
-                offset_x : offset_x + length_x
-            ] = to_write_uncertainty
+            uncert_cpu = uncert_weighted.cpu().numpy()
+
+        # Accumulate in global map
+        for i in range(B):
+            r, c = r_offs[i].item(), c_offs[i].item()
+            h, w = h_actuals[i].item(), w_actuals[i].item()
+
+            # Slice valid regions
+            valid_probs = probs_cpu[i, :, :h, :w]
+            valid_weights = weights_cpu[i, :, :h, :w]
             
+            out_probs[:, r:r+h, c:c+w] += valid_probs
+            count_map[:, r:r+h, c:c+w] += valid_weights
+
+            if is_ensemble:
+                valid_uncert = uncert_cpu[i, :, :h, :w]
+                out_uncert[:, r:r+h, c:c+w] += valid_uncert
+
+    # --- 8. Normalization ---
+    mask_zero = (count_map == 0)
+    count_map[mask_zero] = 1.0 # Prevent div/0
+    
+    out_probs /= count_map
+    out_probs[mask_zero] = nodata
+
     if is_ensemble:
-        return output_probs, output_uncertainty
-    else:
-        return output_probs
+        out_uncert /= count_map
+        out_uncert[mask_zero] = nodata
+        return out_probs, out_uncert
+    
+    return out_probs