add n bands support & rf implementation

slurp/eomultiprocessing/slurp_executor.py

@@ -352,19 +352,39 @@ def mp_n_to_m_images(
     )
 
 
-def _find_spatial_axes(arr: np.ndarray, h: int, w: int):
+def _find_spatial_axes(arr: np.ndarray, h: int, w: int) -> tuple[int, int]:
     """
-    Find which axes correspond to spatial dimensions (H, W)
-    regardless of axis order.
-
-    Returns:
-        (h_axis, w_axis)
+    Find spatial axes (H, W) inside an arbitrary tensor layout.
+
+    The function does NOT assume spatial axes are the last dimensions.
+
+    Examples
+    --------
+    (H, W)           -> (0, 1)
+    (C, H, W)        -> (1, 2)
+    (H, W, C)        -> (0, 1)
+    (T, C, H, W)     -> (2, 3)
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Input tensor.
+    h : int
+        Expected spatial height.
+    w : int
+        Expected spatial width.
+
+    Returns
+    -------
+    tuple[int, int]
+        Indices of spatial axes.
     """
 
     matches = []
 
     for i, size_i in enumerate(arr.shape):
         for j, size_j in enumerate(arr.shape):
+
             if i == j:
                 continue
 
@@ -373,12 +393,13 @@ def _find_spatial_axes(arr: np.ndarray, h: int, w: int):
 
     if not matches:
         raise ValueError(
-            f"Cannot find spatial axes matching ({h},{w}) "
+            f"Cannot find spatial axes matching ({h}, {w}) "
             f"in array shape {arr.shape}"
         )
 
-    # take first valid match
-    return matches[0]
+    # Prefer the last matching pair
+    # because spatial axes are usually deeper
+    return matches[-1]
 
 
 # ============================================================

slurp/masks/watermask.py

@@ -241,10 +241,7 @@ def build_samples(
     valid_stack: np.ndarray,
     mask_hand: np.ndarray,
     mask_pekel: np.ndarray,
-    vhr1: np.ndarray,
-    vhr2: np.ndarray,
-    vhr3: np.ndarray,
-    vhr4: np.ndarray,
+    vhr: np.ndarray,
     ndvi: np.ndarray,
     ndwi: np.ndarray,
     ndwi_threshold: float,
@@ -299,7 +296,6 @@ def build_samples(
     """
     if aux_inputs is None:
         aux_inputs = []
-
     # ---- validity mask ----
     validity_mask = valid_stack == 0
 
@@ -359,7 +355,13 @@ def build_samples(
     cols = np.concatenate((cols_pekel, cols_hand))
 
     # ---- stack features ----
-    base_features = [mask_pekel, vhr1, vhr2, vhr3, vhr4, ndvi, ndwi]
+    vhr_features = [vhr[i] for i in range(vhr.shape[0])]
+    base_features = [
+        mask_pekel,
+        *vhr_features,
+        ndvi,
+        ndwi,
+    ]
     features = base_features + list(aux_inputs)
     im_stack = np.stack(features, axis=0)
 
@@ -373,40 +375,15 @@ def build_samples(
 
 def rf_prediction(
     valid_stack: np.ndarray,
-    vhr1: np.ndarray,
-    vhr2: np.ndarray,
-    vhr3: np.ndarray,
-    vhr4: np.ndarray,
+    vhr: np.ndarray,
     ndvi: np.ndarray,
     ndwi: np.ndarray,
     aux_inputs=None,
     classifier=None,
     debug: bool = False,
 ):
     """
-    Random Forest prediction
-
-    Parameters
-    ----------
-    valid_stack : np.ndarray
-        Validity mask (1, H, W)
-    vhr : np.ndarray
-        VHR features
-    ndvi : np.ndarray
-        NDVI layer(s)
-    ndwi : np.ndarray
-        NDWI layer(s)
-    classifier : sklearn-like estimator
-        Trained classifier
-    debug : bool
-        Enable memory debug logs
-    aux_inputs : np.ndarray
-        Additional feature layers
-
-    Returns
-    -------
-    np.ndarray
-        Predicted mask (uint8)
+    Random Forest prediction on multiband VHR tensor.
     """
     # -------------------------------------------------
     # AUX INPUT NORMALIZATION
@@ -416,33 +393,54 @@ def rf_prediction(
 
     if isinstance(aux_inputs, np.ndarray):
         aux_inputs = [aux_inputs]
-    # ---- build feature stack ----------------------------------------------
-    feature_layers = [
-        vhr1,
-        vhr2,
-        vhr3,
-        vhr4,
-        ndvi,
-        ndwi,
-    ] + aux_inputs
 
-    # concatenate exactly like legacy algo
+    # -------------------------------------------------
+    # MULTIBAND FEATURE EXTRACTION
+    # -------------------------------------------------
+    vhr_features = [
+        vhr[i]
+        for i in range(vhr.shape[0])
+    ]
+
+    # -------------------------------------------------
+    # BUILD FEATURE STACK
+    # -------------------------------------------------
+    feature_layers = (
+        vhr_features
+        + [ndvi, ndwi]
+        + list(aux_inputs)
+    )
+
     im_stack = np.concatenate(
         [layer[np.newaxis, ...] for layer in feature_layers],
         axis=0,
     )
-    # valid_stack shape expected: (1, H, W)
+
+    # -------------------------------------------------
+    # VALID MASK
+    # -------------------------------------------------
     valid_mask = np.logical_not(valid_stack)
 
-    # ---- reshape for sklearn ----------------------------------------------
-    buffer_to_predict = np.transpose(im_stack[:, valid_mask])
+    # -------------------------------------------------
+    # SKLEARN BUFFER
+    # -------------------------------------------------
+    buffer_to_predict = np.transpose(
+        im_stack[:, valid_mask]
+    )
 
-    prediction = np.zeros(im_stack.shape[1:], dtype=np.uint8)
+    prediction = np.zeros(
+        im_stack.shape[1:],
+        dtype=np.uint8,
+    )
 
     if buffer_to_predict.shape[0] > 0:
-        prediction[valid_mask] = classifier.predict(buffer_to_predict)
+        prediction[valid_mask] = classifier.predict(
+            buffer_to_predict
+        )
 
-    # ---- debug -------------------------------------------------------------
+    # -------------------------------------------------
+    # DEBUG
+    # -------------------------------------------------
     utils.display_mem_usage(
         debug,
         f"RF Prediction on buffer "
@@ -905,10 +903,7 @@ def nominal_case_predict(
         valid_stack[0][0],
         mask_hand[0],
         local_mask_pekel,
-        vhr[0][0],
-        vhr[0][1],
-        vhr[0][2],
-        vhr[0][3],
+        vhr[0],
         ndvi[0][0],
         ndwi[0][0],
     ]
@@ -963,10 +958,7 @@ def nominal_case_predict(
     # -- Predict full tile -- #
     input_for_prediction = [
         valid_stack[0][0],
-        vhr[0][0],
-        vhr[0][1],
-        vhr[0][2],
-        vhr[0][3],
+        vhr[0],          # <-- tensor multibande complet
         ndvi[0][0],
         ndwi[0][0],
     ]

slurp/tools/random_forest_utils.py

@@ -31,21 +31,58 @@
 logger = logging.getLogger("slurp")
 
 
+import numpy as np
+
+
 def print_feature_importance(classifier, layers):
-    """Compute feature importance."""
-    feature_names = ["R", "G", "B", "NIR", "NDVI", "NDWI"] + layers
+    """
+    Print feature importance ranking.
+
+    Features are assumed to be ordered as:
+        B1..BN, NDVI, NDWI, auxiliary layers.
+    """
+
+    n_features = len(classifier.feature_importances_)
+    n_aux_layers = len(layers)
+
+    # NDVI + NDWI
+    n_indices = 2
+
+    n_bands = n_features - n_aux_layers - n_indices
+
+    if n_bands <= 0:
+        raise ValueError(
+            f"Unable to infer spectral band count: "
+            f"{n_features=} {n_aux_layers=}."
+        )
+
+    feature_names = (
+        [f"B{i}" for i in range(1, n_bands + 1)]
+        + ["NDVI", "NDWI"]
+        + list(layers)
+    )
+
+    if len(feature_names) != n_features:
+        raise ValueError(
+            f"Feature count mismatch: "
+            f"{len(feature_names)} names generated for "
+            f"{n_features} model features."
+        )
 
     importances = classifier.feature_importances_
     indices = np.argsort(importances)[::-1]
 
     std = np.std(
-        [tree.feature_importances_ for tree in classifier.estimators_], axis=0
+        [tree.feature_importances_ for tree in classifier.estimators_],
+        axis=0,
     )
 
     logger.info("Feature ranking:")
     for idx in indices:
         logger.info(
-            f" {feature_names[idx]:4s} ({importances[idx]:f}) (std={std[idx]:f})"
+            f"{feature_names[idx]:<10s} "
+            f"({importances[idx]:.6f}) "
+            f"(std={std[idx]:.6f})"
         )
 
 

tests/config_test_n_bands.json

@@ -0,0 +1,6 @@
+{
+  "main_config": "tests/main_n_bands_config.json",
+  "features_test_img": "/work/CAMPUS/etudes/Masques_CO3D/Validation/Input-Data/uc_toulouse_with_n_bands/xt_Toulouse_with_n_bands.tif",
+  "ref_dir": "/work/CAMPUS/etudes/Masques_CO3D/Validation/Input-Data/uc_toulouse_with_n_bands/ref",
+  "valid_stack": "/work/CAMPUS/etudes/Masques_CO3D/Validation/Input-Data/uc_toulouse_with_n_bands/ref/valid_stack.tif"
+}