Update load.py

load.py

@@ -121,8 +121,8 @@ def compiled_model(
 def predict_large(
     image: np.ndarray,
     model: nn.Module,
-    chunk_size: int = 512,
-    overlap: int = 256,
+    chunk_size: int = 512,     
+    overlap: int = None,       
     batch_size: int = 1,
     device: str = "cpu",
     merge_clouds: bool = False,
@@ -136,7 +136,7 @@ def predict_large(
         image: Input image (C, H, W) in reflectance [0, 1]
         model: Loaded model from compiled_model()
         chunk_size: Size of inference tiles (default: 512)
-        overlap: Overlap between tiles (default: 256)
+        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
@@ -150,6 +150,7 @@ def predict_large(
     model.eval()
     model.to(device)
     
+    # Get merge_clouds from model if not specified
     if not hasattr(model, 'merge_clouds'):
         model.merge_clouds = merge_clouds
     else:
@@ -157,17 +158,23 @@ def predict_large(
     
     C, H, W = image.shape
     
+    # Set default overlap if not specified
+    if overlap is None:
+        overlap = chunk_size // 2
+    
+    # Direct inference if image fits within chunk_size
     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]
-                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 (thin+thick)
+                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)
@@ -177,9 +184,11 @@ def predict_large(
         
         return pred
     
+    # Sliding window inference for larger images
     step = chunk_size - overlap
     half_tile = chunk_size // 2
     
+    # Pad image to ensure tiles cover entire image
     image_padded = np.pad(
         image,
         ((0, 0), (half_tile, half_tile + chunk_size), (half_tile, half_tile + chunk_size)),
@@ -188,49 +197,61 @@ def predict_large(
     
     _, H_pad, W_pad = image_padded.shape
     
+    # Initialize accumulation buffers
     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)
     
+    # Create 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
             ])
             
+            # 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
     probs_final = probs_final[:, half_tile:half_tile + H, half_tile:half_tile + 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]
-        probs_merged[1] = probs_final[1] + probs_final[2]
-        probs_merged[2] = probs_final[3]
+        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
         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 requested
     if apply_rules:
         pred = apply_physical_rules(pred, image, merge_clouds)