Update app.py

app.py

@@ -70,16 +70,24 @@ def normalise_sar(sar5: np.ndarray) -> np.ndarray:
     stds  = np.array(cfg.SAR_STDS,  dtype=np.float32).reshape(5, 1, 1)
     return (sar5 - means) / (stds + 1e-8)
 
-def load_jrc_channel(img_path: Path, shape: tuple) -> np.ndarray:
+def load_jrc_channel(img_path: Path, shape: tuple, raw: np.ndarray = None) -> np.ndarray:
+    # Priority 1: embedded as band 7 in the uploaded TIF
+    if raw is not None and raw.shape[0] >= 7:
+        mask = raw[6].copy()
+        mask[mask > 1] = 0.0
+        return mask[np.newaxis]
+    
+    # Priority 2: separate *_waterbody.tif on disk
     stem = img_path.stem.replace("_image", "")
     mask_path = cfg.JRC_DIR / f"{stem}_waterbody.tif"
     if mask_path.exists():
         with rasterio.open(mask_path) as src:
             mask = src.read(1).astype(np.float32)
         mask[mask > 1] = 0.0
-    else:
-        mask = np.zeros(shape, dtype=np.float32)
-    return mask[np.newaxis]
+        return mask[np.newaxis]
+    
+    # Fallback: zeros
+    return np.zeros((1, *shape), dtype=np.float32)
 
 # ═══════════════════════════════════════════════════════════════════════════════
 #  MODEL ARCHITECTURE
@@ -295,11 +303,13 @@ def sliding_window_predict(model, image_tensor, crop=224, stride=112, device="cu
 
 def predict_image(model, img_path: Path, device="cuda"):
     raw = read_tif(img_path)
-    opt = normalise_optical(raw); sar = normalise_sar(compute_sar_features(raw))
-    jrc = load_jrc_channel(img_path, raw.shape[1:])
+    opt = normalise_optical(raw[:6])   # always use first 6 bands for optical
+    sar = normalise_sar(compute_sar_features(raw[:6]))
+    jrc = load_jrc_channel(img_path, raw.shape[1:], raw)  # pass raw for embedded check
     img = torch.from_numpy(np.concatenate([opt, sar, jrc], axis=0)).float()
     return sliding_window_predict(model, img, cfg.IMG_SIZE, cfg.INFER_STRIDE, device).cpu().numpy()
 
+
 def mask_to_rle(mask: np.ndarray) -> str:
     flat = mask.flatten(order="F").astype(np.uint8)
     if flat.max() == 0: return "0 0"