Update load.py

load.py

@@ -121,7 +121,6 @@ def compiled_model(
     
     return model
 
-
 def predict_large(
     image: np.ndarray,
     model: nn.Module,
@@ -131,24 +130,29 @@ def predict_large(
     device: str = "cpu",
     merge_clouds: bool = False,
     apply_rules: bool = False,
+    max_direct_size: int = 1024,  # Safe for 2GB GPU
     **kwargs
 ) -> np.ndarray:
     """
     Predict on images of any size.
     
     Strategy:
-    - Images ≤ 2048px in any dimension: direct inference (no tiling)
-    - Images > 2048px: sliding window with specified chunk_size
+    - 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: Size of inference tiles for large images (default: 512)
+        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)
@@ -168,10 +172,9 @@ def predict_large(
     if overlap is None:
         overlap = chunk_size // 2
     
-    # === STRATEGY: Use direct inference for images ≤ 2048px ===
-    # This avoids tiling artifacts on small/medium images
-    if max(H, W) <= 2048:
-        # Direct inference - no tiling
+    # === 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)
@@ -179,9 +182,9 @@ def predict_large(
             if merge_clouds:
                 probs = torch.softmax(logits, dim=1)
                 probs_merged = torch.zeros(1, 3, H, W, device=device)
-                probs_merged[:, 0] = probs[:, 0]
-                probs_merged[:, 1] = probs[:, 1] + probs[:, 2]
-                probs_merged[:, 2] = probs[:, 3]
+                probs_merged[:, 0] = probs[:, 0]  # Clear
+                probs_merged[:, 1] = probs[:, 1] + probs[:, 2]  # Cloud
+                probs_merged[:, 2] = probs[:, 3]  # Shadow
                 pred = probs_merged.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
             else:
                 pred = logits.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
@@ -191,7 +194,7 @@ def predict_large(
         
         return pred
     
-    # === SLIDING WINDOW FOR LARGE IMAGES (> 2048px) ===
+    # === SLIDING WINDOW FOR LARGE IMAGES ===
     
     step = chunk_size - overlap