Update load.py

load.py

@@ -7,16 +7,11 @@ import torch
 import torch.nn as nn
 import numpy as np
 from pathlib import Path
-from typing import Tuple, Optional
 import pytorch_lightning as pl
 import segmentation_models_pytorch as smp
 from tqdm import tqdm
 
 
-# ============================================================================
-# MODEL DEFINITION (copied from your model.py)
-# ============================================================================
-
 class MSSSegmentationModel(pl.LightningModule):
     """UNet para cloud segmentation en MSS."""
     
@@ -45,12 +40,8 @@ class MSSSegmentationModel(pl.LightningModule):
         return self.model(x)
 
 
-# ============================================================================
-# INFERENCE UTILITIES
-# ============================================================================
-
 def get_spline_window(size: int, power: int = 2) -> np.ndarray:
-    """Generate Hann window for smooth blending."""
+    """Ventana Hann 2D para blending suave."""
     intersection = np.hanning(size)
     window_2d = np.outer(intersection, intersection)
     return (window_2d ** power).astype(np.float32)
@@ -60,42 +51,28 @@ def apply_physical_rules(
     pred: np.ndarray,
     image: np.ndarray,
     merge_clouds: bool = False,
-    saturation_threshold: float = 0.35,
 ) -> np.ndarray:
-    """
-    Apply physical rules for better cloud detection.
+    """Regla física para nubes gruesas saturadas."""
+    saturation_threshold = 0.35
     
-    Args:
-        pred: Predicted classes (H, W)
-        image: Input image (4, H, W) in reflectance [0, 1]
-        merge_clouds: If True, merge thin+thick into single cloud class
-        saturation_threshold: Threshold for detecting saturated bright clouds
-    """
     pred = pred.copy()
     
-    # Mask nodata pixels
     nodata_mask = np.all(image == 0, axis=0)
-    pred[nodata_mask] = 0
     
-    # Detect very bright pixels (likely thick clouds)
     bright_b0 = image[0] > saturation_threshold
     bright_b1 = image[1] > saturation_threshold * 0.80
     saturated_mask = bright_b0 & bright_b1
     
     if merge_clouds:
-        # Set to cloud (1)
         pred[saturated_mask] = 1
     else:
-        # Set to thick cloud (2)
         pred[saturated_mask] = 2
     
+    pred[nodata_mask] = 0
+    
     return pred
 
 
-# ============================================================================
-# MLSTAC-COMPATIBLE FUNCTIONS
-# ============================================================================
-
 def compiled_model(
     model_dir: Path,
     stac_item=None,
@@ -110,20 +87,18 @@ def compiled_model(
         model_dir: Directory containing the .ckpt file
         stac_item: STAC item metadata (optional)
         device: 'cpu' or 'cuda'
-        merge_clouds: If True, output will have 3 classes (clear, cloud, shadow)
-                     If False, output will have 4 classes (clear, thin, thick, shadow)
+        merge_clouds: If True, output 3 classes (clear, cloud, shadow)
+                     If False, output 4 classes (clear, thin, thick, shadow)
     
     Returns:
         Loaded model in eval mode
     """
-    # Find checkpoint file
     ckpt_files = list(model_dir.glob("*.ckpt"))
     if not ckpt_files:
         raise FileNotFoundError(f"No .ckpt file found in {model_dir}")
     
     ckpt_path = ckpt_files[0]
     
-    # Load model
     model = MSSSegmentationModel.load_from_checkpoint(
         ckpt_path,
         map_location=device
@@ -131,11 +106,9 @@ def compiled_model(
     model.eval()
     model.to(device)
     
-    # Disable gradients
     for param in model.parameters():
         param.requires_grad = False
     
-    # Store merge_clouds flag for predict_large
     model.merge_clouds = merge_clouds
     
     print(f"✅ Model loaded from {ckpt_path.name}")
@@ -152,9 +125,8 @@ def predict_large(
     overlap: int = 256,
     batch_size: int = 1,
     device: str = "cpu",
-    nodata: float = 0.0,
-    apply_rules: bool = True,
-    saturation_threshold: float = 0.35,
+    merge_clouds: bool = False,
+    apply_rules: bool = False,
     **kwargs
 ) -> np.ndarray:
     """
@@ -163,53 +135,51 @@ def predict_large(
     Args:
         image: Input image (C, H, W) in reflectance [0, 1]
         model: Loaded model from compiled_model()
-        chunk_size: Size of inference tiles (default: 1024)
-        overlap: Overlap between tiles for smooth blending (default: 256)
-        batch_size: Number of tiles to process in parallel (default: 1)
+        chunk_size: Size of inference tiles (default: 512)
+        overlap: Overlap between tiles (default: 256)
+        batch_size: Tiles per batch (default: 1)
         device: 'cpu' or 'cuda'
-        nodata: Value representing no-data pixels
-        apply_rules: Whether to apply physical rules post-processing
-        saturation_threshold: Threshold for detecting bright clouds
+        merge_clouds: If True, merge thin+thick into single cloud class
+        apply_rules: If True, apply physical rules for bright clouds (default: False)
     
     Returns:
-        Predicted class labels (H, W) with shape matching input
+        Predicted class labels (H, W)
         - If merge_clouds=False: 0=clear, 1=thin, 2=thick, 3=shadow
         - If merge_clouds=True: 0=clear, 1=cloud, 2=shadow
     """
     model.eval()
     model.to(device)
     
-    merge_clouds = getattr(model, 'merge_clouds', False)
+    if not hasattr(model, 'merge_clouds'):
+        model.merge_clouds = merge_clouds
+    else:
+        merge_clouds = model.merge_clouds
     
     C, H, W = image.shape
     
-    # Direct inference for small images
     if H <= chunk_size and W <= chunk_size:
         with torch.no_grad():
             img_tensor = torch.from_numpy(image).unsqueeze(0).float().to(device)
             logits = model(img_tensor)
             
             if merge_clouds:
-                # Merge thin(1) + thick(2) probabilities
                 probs = torch.softmax(logits, dim=1)
                 probs_merged = torch.zeros(1, 3, H, W, device=device)
-                probs_merged[:, 0] = probs[:, 0]  # clear
-                probs_merged[:, 1] = probs[:, 1] + probs[:, 2]  # cloud
-                probs_merged[:, 2] = probs[:, 3]  # shadow
+                probs_merged[:, 0] = probs[:, 0]
+                probs_merged[:, 1] = probs[:, 1] + probs[:, 2]
+                probs_merged[:, 2] = probs[:, 3]
                 pred = probs_merged.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
             else:
                 pred = logits.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
         
         if apply_rules:
-            pred = apply_physical_rules(pred, image, merge_clouds, saturation_threshold)
+            pred = apply_physical_rules(pred, image, merge_clouds)
         
         return pred
     
-    # Sliding window for large images
     step = chunk_size - overlap
     half_tile = chunk_size // 2
     
-    # Pad image
     image_padded = np.pad(
         image,
         ((0, 0), (half_tile, half_tile + chunk_size), (half_tile, half_tile + chunk_size)),
@@ -218,82 +188,60 @@ def predict_large(
     
     _, H_pad, W_pad = image_padded.shape
     
-    # Initialize accumulators - ALWAYS 4 classes, merge at the end if needed
     num_classes = 4
     probs_sum = np.zeros((num_classes, H_pad, W_pad), dtype=np.float32)
     weight_sum = np.zeros((H_pad, W_pad), dtype=np.float32)
     
-    # Blending window
     window = get_spline_window(chunk_size, power=2)
     
-    # Generate tile coordinates
     coords = [
         (r, c)
         for r in range(0, H_pad - chunk_size + 1, step)
         for c in range(0, W_pad - chunk_size + 1, step)
     ]
     
-    # Process tiles in batches
     with torch.no_grad():
         for i in tqdm(range(0, len(coords), batch_size), desc="  Tiles", leave=False, disable=True):
             batch_coords = coords[i:i + batch_size]
             
-            # Extract tiles
             tiles = np.stack([
                 image_padded[:, r:r + chunk_size, c:c + chunk_size]
                 for r, c in batch_coords
             ])
             
-            # Inference
             tiles_tensor = torch.from_numpy(tiles).float().to(device)
             logits = model(tiles_tensor)
             probs = torch.softmax(logits, dim=1).cpu().numpy()
             
-            # Accumulate with blending - ALWAYS accumulate 4 classes
             for j, (r, c) in enumerate(batch_coords):
                 probs_sum[:, r:r + chunk_size, c:c + chunk_size] += probs[j] * window
                 weight_sum[r:r + chunk_size, c:c + chunk_size] += window
     
-    # Normalize
     weight_sum = np.maximum(weight_sum, 1e-8)
     probs_final = probs_sum / weight_sum
     
-    # Crop to original size
     probs_final = probs_final[:, half_tile:half_tile + H, half_tile:half_tile + W]
     
-    # Merge classes if requested - AFTER normalization
     if merge_clouds:
         probs_merged = np.zeros((3, H, W), dtype=np.float32)
-        probs_merged[0] = probs_final[0]  # clear
-        probs_merged[1] = probs_final[1] + probs_final[2]  # cloud = thin + thick
-        probs_merged[2] = probs_final[3]  # shadow
+        probs_merged[0] = probs_final[0]
+        probs_merged[1] = probs_final[1] + probs_final[2]
+        probs_merged[2] = probs_final[3]
         pred = np.argmax(probs_merged, axis=0).astype(np.uint8)
     else:
         pred = np.argmax(probs_final, axis=0).astype(np.uint8)
     
-    # Apply physical rules
     if apply_rules:
-        pred = apply_physical_rules(pred, image, merge_clouds, saturation_threshold)
+        pred = apply_physical_rules(pred, image, merge_clouds)
     
     return pred
 
 
-# ============================================================================
-# OPTIONAL: EXAMPLE DATA AND VISUALIZATION
-# ============================================================================
-
 def example_data(model_dir: Path, **kwargs):
-    """
-    Load example data for testing (optional function).
-    
-    Returns:
-        Example MSS image as numpy array (4, H, W)
-    """
-    # This is optional - you can provide a small example .npy file
+    """Load example data for testing."""
     example_path = model_dir / "example_mss.npy"
     
     if not example_path.exists():
-        # Return synthetic data if no example file
         print("⚠️  No example data found, generating synthetic")
         return np.random.rand(4, 512, 512).astype(np.float32) * 0.5
     
@@ -307,15 +255,7 @@ def display_results(
     stac_item=None,
     **kwargs
 ):
-    """
-    Display prediction results (optional visualization function).
-    
-    Args:
-        model_dir: Model directory
-        image: Input image (4, H, W)
-        prediction: Predicted classes (H, W)
-        stac_item: STAC metadata
-    """
+    """Display prediction results."""
     try:
         import matplotlib.pyplot as plt
         from matplotlib.colors import ListedColormap
@@ -325,9 +265,8 @@ def display_results(
     
     merge_clouds = prediction.max() <= 2
     
-    # Color maps
     if merge_clouds:
-        colors = ['#2E7D32', '#FFFFFF', '#424242']  # clear, cloud, shadow
+        colors = ['#2E7D32', '#FFFFFF', '#424242']
         labels = ['Clear', 'Cloud', 'Shadow']
     else:
         colors = ['#2E7D32', '#B3E5FC', '#FFFFFF', '#424242']
@@ -335,22 +274,18 @@ def display_results(
     
     cmap = ListedColormap(colors)
     
-    # Plot
     fig, axes = plt.subplots(1, 2, figsize=(12, 5))
     
-    # RGB composite (use bands 1, 0, 2 as RGB approximation)
     rgb = np.stack([image[1], image[0], image[2]], axis=-1)
-    rgb = np.clip(rgb * 3, 0, 1)  # Brighten for visibility
+    rgb = np.clip(rgb * 3, 0, 1)
     axes[0].imshow(rgb)
     axes[0].set_title("MSS RGB Composite")
     axes[0].axis('off')
     
-    # Prediction
     im = axes[1].imshow(prediction, cmap=cmap, vmin=0, vmax=len(labels)-1)
     axes[1].set_title("Cloud Detection")
     axes[1].axis('off')
     
-    # Colorbar
     cbar = plt.colorbar(im, ax=axes[1], ticks=range(len(labels)))
     cbar.ax.set_yticklabels(labels)