Update load.py

load.py

@@ -41,7 +41,7 @@ class MSSSegmentationModel(pl.LightningModule):
 
 
 def get_spline_window(size: int, power: int = 2) -> np.ndarray:
-    """Ventana Hann 2D para blending suave."""
+    """Hann window for smooth blending."""
     intersection = np.hanning(size)
     window_2d = np.outer(intersection, intersection)
     return (window_2d ** power).astype(np.float32)
@@ -52,22 +52,26 @@ def apply_physical_rules(
     image: np.ndarray,
     merge_clouds: bool = False,
 ) -> np.ndarray:
-    """Regla física para nubes gruesas saturadas."""
+    """Apply physical rules for saturated thick clouds."""
     saturation_threshold = 0.35
     
     pred = pred.copy()
     
+    # Nodata mask
     nodata_mask = np.all(image == 0, axis=0)
     
+    # Saturated clouds (high values in visible bands)
     bright_b0 = image[0] > saturation_threshold
     bright_b1 = image[1] > saturation_threshold * 0.80
     saturated_mask = bright_b0 & bright_b1
     
+    # Assign thick cloud class
     if merge_clouds:
-        pred[saturated_mask] = 1
+        pred[saturated_mask] = 1  # Cloud (merged)
     else:
-        pred[saturated_mask] = 2
+        pred[saturated_mask] = 2  # Thick cloud
     
+    # Set nodata to clear
     pred[nodata_mask] = 0
     
     return pred
@@ -121,8 +125,8 @@ def compiled_model(
 def predict_large(
     image: np.ndarray,
     model: nn.Module,
-    chunk_size: int = 512,     
-    overlap: int = None,       
+    chunk_size: int = 512,
+    overlap: int = None,
     batch_size: int = 1,
     device: str = "cpu",
     merge_clouds: bool = False,
@@ -130,7 +134,7 @@ def predict_large(
     **kwargs
 ) -> np.ndarray:
     """
-    Predict on large images using sliding window with overlap blending.
+    Predict on images of any size using sliding window with smooth blending.
     
     Args:
         image: Input image (C, H, W) in reflectance [0, 1]
@@ -140,7 +144,7 @@ def predict_large(
         batch_size: Tiles per batch (default: 1)
         device: 'cpu' or 'cuda'
         merge_clouds: If True, merge thin+thick into single cloud class
-        apply_rules: If True, apply physical rules for bright clouds (default: False)
+        apply_rules: If True, apply physical rules for bright clouds
     
     Returns:
         Predicted class labels (H, W)
@@ -150,7 +154,7 @@ def predict_large(
     model.eval()
     model.to(device)
     
-    # Get merge_clouds from model if not specified
+    # Get merge_clouds setting from model if available
     if not hasattr(model, 'merge_clouds'):
         model.merge_clouds = merge_clouds
     else:
@@ -158,23 +162,22 @@ def predict_large(
     
     C, H, W = image.shape
     
-    # Set default overlap if not specified
+    # Set default overlap
     if overlap is None:
         overlap = chunk_size // 2
     
-    # Direct inference if image fits within chunk_size
+    # 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+thick clouds into single cloud class
                 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 (thin+thick)
-                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)
@@ -184,14 +187,23 @@ def predict_large(
         
         return pred
     
-    # Sliding window inference for larger images
+    # === SLIDING WINDOW FOR LARGER IMAGES ===
+    
+    # Calculate padding needed to make image divisible by step
     step = chunk_size - overlap
-    half_tile = chunk_size // 2
     
-    # Pad image to ensure tiles cover entire image
+    # Padding to ensure tiles cover the entire image
+    pad_h = (step - (H % step)) % step
+    pad_w = (step - (W % step)) % step
+    
+    # Add extra overlap padding on all sides for smooth edges
+    pad_h += overlap
+    pad_w += overlap
+    
+    # Pad image
     image_padded = np.pad(
         image,
-        ((0, 0), (half_tile, half_tile + chunk_size), (half_tile, half_tile + chunk_size)),
+        ((0, 0), (overlap // 2, pad_h - overlap // 2), (overlap // 2, pad_w - overlap // 2)),
         mode="reflect"
     )
     
@@ -206,15 +218,14 @@ def predict_large(
     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)
-    ]
+    coords = []
+    for r in range(0, H_pad - chunk_size + 1, step):
+        for c in range(0, W_pad - chunk_size + 1, step):
+            coords.append((r, c))
     
     # Process tiles in batches
     with torch.no_grad():
-        for i in tqdm(range(0, len(coords), batch_size), desc="  Tiles", leave=False, disable=True):
+        for i in range(0, len(coords), batch_size):
             batch_coords = coords[i:i + batch_size]
             
             # Extract tiles
@@ -237,16 +248,15 @@ def predict_large(
     weight_sum = np.maximum(weight_sum, 1e-8)
     probs_final = probs_sum / weight_sum
     
-    # Remove padding
-    probs_final = probs_final[:, half_tile:half_tile + H, half_tile:half_tile + W]
+    # Remove padding to get back to original size
+    probs_final = probs_final[:, overlap // 2:overlap // 2 + H, overlap // 2:overlap // 2 + W]
     
     # Get final prediction
     if merge_clouds:
-        # Merge thin+thick clouds into single cloud class
         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)
@@ -297,16 +307,19 @@ def display_results(
     
     fig, axes = plt.subplots(1, 2, figsize=(12, 5))
     
+    # RGB composite
     rgb = np.stack([image[1], image[0], image[2]], axis=-1)
     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)