Update load.py

load.py

@@ -125,29 +125,30 @@ def compiled_model(
 def predict_large(
     image: np.ndarray,
     model: nn.Module,
-    chunk_size: int = None,  # None = auto
+    chunk_size: int = 512,
     overlap: int = None,
     batch_size: int = 1,
     device: str = "cpu",
     merge_clouds: bool = False,
     apply_rules: bool = False,
-    max_direct_size: int = 4096,  # Max size for direct inference
     **kwargs
 ) -> np.ndarray:
     """
     Predict on images of any size.
-    Automatically uses direct inference for small/medium images to avoid tiling artifacts.
+    
+    Strategy:
+    - Images ≤ 2048px in any dimension: direct inference (no tiling)
+    - Images > 2048px: sliding window with specified chunk_size
     
     Args:
         image: Input image (C, H, W) in reflectance [0, 1]
         model: Loaded model from compiled_model()
-        chunk_size: Size of inference tiles (None = auto-detect)
-        overlap: Overlap between tiles (None = auto, chunk_size // 2)
+        chunk_size: Size of inference tiles 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: Maximum dimension for direct inference (default: 4096)
     
     Returns:
         Predicted class labels (H, W)
@@ -162,24 +163,15 @@ def predict_large(
         merge_clouds = model.merge_clouds
     
     C, H, W = image.shape
-    max_dim = max(H, W)
-    
-    # === AUTO CHUNK SIZE ===
-    if chunk_size is None:
-        if max_dim <= 1024:
-            chunk_size = max_dim  # Process entire image
-        elif max_dim <= 2048:
-            chunk_size = 1024
-        else:
-            chunk_size = 512
     
     # Set default overlap
     if overlap is None:
         overlap = chunk_size // 2
     
-    # === DIRECT INFERENCE (NO TILING) ===
-    # Use direct inference if image fits in single chunk or is small enough
-    if max_dim <= max_direct_size or (H <= chunk_size and W <= chunk_size):
+    # === STRATEGY: Use direct inference for images ≤ 2048px ===
+    # This avoids tiling artifacts on small/medium images
+    if max(H, W) <= 2048:
+        # Direct inference - no tiling
         with torch.no_grad():
             img_tensor = torch.from_numpy(image).unsqueeze(0).float().to(device)
             logits = model(img_tensor)
@@ -199,8 +191,7 @@ def predict_large(
         
         return pred
     
-    # === SLIDING WINDOW FOR VERY LARGE IMAGES ===
-    print(f"  Using sliding window: {H}x{W} -> chunks={chunk_size}, overlap={overlap}")
+    # === SLIDING WINDOW FOR LARGE IMAGES (> 2048px) ===
     
     step = chunk_size - overlap
     
@@ -230,7 +221,7 @@ def predict_large(
     # Create blending window
     window = get_spline_window(chunk_size, power=2)
     
-    # Generate tile coordinates - ensure full coverage
+    # 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):