Update load.py

load.py

@@ -134,7 +134,9 @@ def predict_large(
     **kwargs
 ) -> np.ndarray:
     """
-    Predict on images of any size using sliding window with smooth blending.
+    Predict on images of any size.
+    - Small images (≤2048px): direct inference without tiling
+    - Large images (>2048px): sliding window with smooth blending
     
     Args:
         image: Input image (C, H, W) in reflectance [0, 1]
@@ -166,8 +168,9 @@ def predict_large(
     if overlap is None:
         overlap = chunk_size // 2
     
-    # Direct inference for small images
-    if H <= chunk_size and W <= chunk_size:
+    # === DIRECT INFERENCE FOR SMALL/MEDIUM IMAGES ===
+    # Process directly without tiling to avoid artifacts
+    if max(H, W) <= 2048:
         with torch.no_grad():
             img_tensor = torch.from_numpy(image).unsqueeze(0).float().to(device)
             logits = model(img_tensor)
@@ -187,23 +190,16 @@ def predict_large(
         
         return pred
     
-    # === SLIDING WINDOW FOR LARGER IMAGES ===
+    # === SLIDING WINDOW FOR LARGE IMAGES (>2048px) ===
     
-    # Calculate padding needed to make image divisible by step
     step = chunk_size - overlap
     
-    # Padding to ensure tiles cover the entire image
-    pad_h = (step - (H % step)) % step
-    pad_w = (step - (W % step)) % step
+    # Symmetric padding: overlap on each side
+    pad_size = overlap
     
-    # 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), (overlap // 2, pad_h - overlap // 2), (overlap // 2, pad_w - overlap // 2)),
+        ((0, 0), (pad_size, pad_size), (pad_size, pad_size)),
         mode="reflect"
     )
     
@@ -219,9 +215,28 @@ def predict_large(
     
     # 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):
+    r = 0
+    while r <= H_pad - chunk_size:
+        c = 0
+        while c <= W_pad - chunk_size:
             coords.append((r, c))
+            c += step
+        # Ensure we cover the right edge
+        if c - step + chunk_size < W_pad:
+            coords.append((r, W_pad - chunk_size))
+        r += step
+    
+    # Ensure we cover the bottom edge
+    if r - step + chunk_size < H_pad:
+        c = 0
+        while c <= W_pad - chunk_size:
+            coords.append((H_pad - chunk_size, c))
+            c += step
+        if c - step + chunk_size < W_pad:
+            coords.append((H_pad - chunk_size, W_pad - chunk_size))
+    
+    # Remove duplicates
+    coords = list(set(coords))
     
     # Process tiles in batches
     with torch.no_grad():
@@ -248,8 +263,8 @@ def predict_large(
     weight_sum = np.maximum(weight_sum, 1e-8)
     probs_final = probs_sum / weight_sum
     
-    # Remove padding to get back to original size
-    probs_final = probs_final[:, overlap // 2:overlap // 2 + H, overlap // 2:overlap // 2 + W]
+    # Remove symmetric padding
+    probs_final = probs_final[:, pad_size:pad_size + H, pad_size:pad_size + W]
     
     # Get final prediction
     if merge_clouds: