Update load.py

load.py

@@ -1,8 +1,3 @@
-"""
-Load and inference functions for MSS Cloud Detection Model
-Compatible with mlstac package
-"""
-
 import torch
 import torch.nn as nn
 import numpy as np
@@ -130,56 +125,37 @@ def predict_large(
     device: str = "cpu",
     merge_clouds: bool = False,
     apply_rules: bool = False,
-    max_direct_size: int = 1024,  # Safe for 2GB GPU
+    max_direct_size: int = 1024,
     **kwargs
 ) -> np.ndarray:
     """
     Predict on images of any size.
     
-    Strategy:
-    - Small images (≤ max_direct_size): direct inference without tiling
-      Examples: 256x256, 512x512, 1024x1024 (safe for 2GB GPU)
-    - Large images (> max_direct_size): sliding window with overlapping tiles
-      Examples: 2048x2048, 5000x5000, 22000x22000
-    
-    Args:
-        image: Input image (C, H, W) in reflectance [0, 1]
-        model: Loaded model from compiled_model()
-        chunk_size: Tile size for large images (default: 512)
-        overlap: Overlap between tiles (default: chunk_size // 2)
-        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
-        max_direct_size: Max dimension for direct inference (default: 1024)
-                        Set to 2048 for GPUs with ≥8GB VRAM
-    
-    Returns:
-        Predicted class labels (H, W)
+    Automatically detects if model has 3 or 4 classes.
     """
     model.eval()
     model.to(device)
     
-    # Get merge_clouds setting from model if available
-    if not hasattr(model, 'merge_clouds'):
-        model.merge_clouds = merge_clouds
-    else:
-        merge_clouds = model.merge_clouds
+    # Detect number of classes in the model
+    num_classes = model.hparams.get('num_classes', 4)
+    is_3class_model = (num_classes == 3)
     
     C, H, W = image.shape
     
-    # Set default overlap
     if overlap is None:
         overlap = chunk_size // 2
     
     # === DIRECT INFERENCE FOR SMALL IMAGES ===
-    # Safe for GPUs with limited VRAM (2-4GB)
     if max(H, W) <= max_direct_size:
         with torch.no_grad():
             img_tensor = torch.from_numpy(image).unsqueeze(0).float().to(device)
             logits = model(img_tensor)
             
-            if merge_clouds:
+            if is_3class_model:
+                # The model already has 3 classes: 0=clear, 1=cloud, 2=shadow
+                pred = logits.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
+            elif merge_clouds:
+                # Model 4 classes → merge to 3
                 probs = torch.softmax(logits, dim=1)
                 probs_merged = torch.zeros(1, 3, H, W, device=device)
                 probs_merged[:, 0] = probs[:, 0]  # Clear
@@ -187,10 +163,11 @@ def predict_large(
                 probs_merged[:, 2] = probs[:, 3]  # Shadow
                 pred = probs_merged.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
             else:
+                # Model 4 classes without merge
                 pred = logits.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
         
         if apply_rules:
-            pred = apply_physical_rules(pred, image, merge_clouds)
+            pred = apply_physical_rules(pred, image, merge_clouds=is_3class_model or merge_clouds)
         
         return pred
     
@@ -198,11 +175,9 @@ def predict_large(
     
     step = chunk_size - overlap
     
-    # Calculate required padding
     pad_h = (step - (H - chunk_size) % step) % step
     pad_w = (step - (W - chunk_size) % step) % step
     
-    # Symmetric padding
     pad_top = pad_h // 2
     pad_bottom = pad_h - pad_top
     pad_left = pad_w // 2
@@ -216,50 +191,45 @@ def predict_large(
     
     _, H_pad, W_pad = image_padded.shape
     
-    # Initialize accumulation buffers
-    num_classes = 4
+    # Buffers according to number of classes
     probs_sum = np.zeros((num_classes, H_pad, W_pad), dtype=np.float32)
     weight_sum = np.zeros((H_pad, W_pad), dtype=np.float32)
     
-    # Create blending window
     window = get_spline_window(chunk_size, power=2)
     
-    # Generate tile coordinates
     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 range(0, len(coords), batch_size):
             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
             ])
             
-            # Run inference
             tiles_tensor = torch.from_numpy(tiles).float().to(device)
             logits = model(tiles_tensor)
             probs = torch.softmax(logits, dim=1).cpu().numpy()
             
-            # Accumulate weighted predictions
             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 by accumulated weights
     weight_sum = np.maximum(weight_sum, 1e-8)
     probs_final = probs_sum / weight_sum
     
-    # Remove padding to restore original size
     probs_final = probs_final[:, pad_top:pad_top + H, pad_left:pad_left + W]
     
-    # Get final prediction
-    if merge_clouds:
+    # Final forecast
+    if is_3class_model:
+        # It already has 3 classes
+        pred = np.argmax(probs_final, axis=0).astype(np.uint8)
+    elif merge_clouds:
+        # Merge 4 → 3
         probs_merged = np.zeros((3, H, W), dtype=np.float32)
         probs_merged[0] = probs_final[0]
         probs_merged[1] = probs_final[1] + probs_final[2]
@@ -268,9 +238,8 @@ def predict_large(
     else:
         pred = np.argmax(probs_final, axis=0).astype(np.uint8)
     
-    # Apply physical rules if requested
     if apply_rules:
-        pred = apply_physical_rules(pred, image, merge_clouds)
+        pred = apply_physical_rules(pred, image, merge_clouds=is_3class_model or merge_clouds)
     
     return pred