Add mlstac loader, weights, examples and model card

load.py

@@ -286,4 +286,122 @@ def predict_large(image, model, **kwargs):
         apply_nodata_mask=kwargs.get("apply_nodata_mask", True),
         return_probs=kwargs.get("return_probs", False),
         verbose=kwargs.get("verbose", False),
-    )
+    )
+
+
+# ============================================================================
+# CHRIS/PROBA-1 preprocessing (the 6 acquisition modes)
+# ============================================================================
+# For each CHRIS mode, which raw band indices (1-based) average into each of
+# the R, G, B, NIR groups. Taken from the inference pipeline.
+BAND_SELECTION = {
+    1: {"B4": [23, 24, 25],
+        "B3": [13, 14, 15],
+        "B2": [4, 5, 6, 7, 8, 9, 10],
+        "B8": [43, 44, 45, 46, 47, 48, 49, 50, 51]},
+    2: {"B4": [10, 11, 12], "B3": [6, 7], "B2": [3, 4], "B8": [17]},
+    3: {"B4": [7, 8],       "B3": [4, 5], "B2": [2],    "B8": [15]},
+    4: {"B4": [4],          "B3": [2],    "B2": [1],    "B8": [18]},
+    5: {"B4": [7, 8],       "B3": [4, 5], "B2": [2],    "B8": [23, 24, 25, 26]},
+    6: {"B2": [1],          "B3": [2],    "B4": [3],    "B8": [4]},
+}
+BAND_KEYS_RGBN = ["B4", "B3", "B2", "B8"]   # R, G, B, NIR
+
+DN_SCALE = 100_000.0
+TOA_SCALE = 10_000.0
+CAP = 5.0   # clip ceiling used during training; keep it identical at inference
+
+
+def _avg_bands(cube, band_list):
+    """Average the given 1-based raw bands into a single (H, W) layer."""
+    idx = [b - 1 for b in band_list]
+    return cube[idx[0]] if len(idx) == 1 else cube[idx].mean(axis=0)
+
+
+def _infer_source_from_name(tif_path):
+    """Guess 'dn' or 'toa' from the file name. Returns None if unclear."""
+    name = Path(tif_path).name.lower()
+    if "toa" in name:
+        return "toa"
+    if "dn" in name:
+        return "dn"
+    return None
+
+
+def build_rgbn(cube, mode_n, source):
+    """Build the (4, H, W) RGBN stack and the nodata mask from a raw cube.
+
+    Args:
+        cube: (bands, H, W) float array read from the CHRIS GeoTIFF.
+        mode_n: CHRIS acquisition mode, 1 to 6.
+        source: 'dn' or 'toa', selects the radiometric scale.
+
+    Returns:
+        (stack, nodata) where stack is (4, H, W) float32 and nodata is (H, W) bool.
+    """
+    sel = BAND_SELECTION.get(int(mode_n))
+    if sel is None:
+        raise ValueError(f"Unsupported CHRIS mode {mode_n}; expected 1-6.")
+
+    needed_max = max(max(sel[k]) for k in BAND_KEYS_RGBN)
+    if cube.shape[0] < needed_max:
+        raise ValueError(
+            f"Cube has {cube.shape[0]} bands but mode {mode_n} needs {needed_max}."
+        )
+
+    # nodata = pixel is 0 in every RGBN group
+    H, W = cube.shape[1:]
+    nodata = np.ones((H, W), dtype=bool)
+    for k in BAND_KEYS_RGBN:
+        nodata &= (_avg_bands(cube, sel[k]) == 0)
+
+    scale = DN_SCALE if source == "dn" else TOA_SCALE
+    layers = [_avg_bands(cube, sel[k]) for k in BAND_KEYS_RGBN]
+    stack = np.stack(layers, axis=0).astype(np.float32) / scale
+    stack = np.clip(stack, 0.0, CAP)
+    stack[:, nodata] = 0.0
+    return stack, nodata
+
+
+def predict_chris(tif_path, model, mode_n, source=None, **kwargs):
+    """Segment a raw CHRIS/PROBA-1 GeoTIFF end to end.
+
+    Reads the cube, builds the RGBN stack for the given mode, runs the
+    ensemble, and restores nodata as 99.
+
+    Args:
+        tif_path: path to the CHRIS GeoTIFF.
+        model: object returned by compiled_model().
+        mode_n: CHRIS acquisition mode, 1 to 6.
+        source: 'dn' or 'toa'. If None, it is guessed from the file name.
+        return_probs: if True, return (num_classes, H, W) probabilities.
+
+    Returns:
+        (H, W) uint8 label map (nodata = 99), or probabilities if requested.
+    """
+    import rasterio as rio
+
+    if source is None:
+        source = _infer_source_from_name(tif_path)
+        if source is None:
+            raise ValueError(
+                "Could not tell DN from TOA by the file name; "
+                "pass source='dn' or source='toa'."
+            )
+
+    with rio.open(tif_path) as src:
+        cube = src.read().astype(np.float32)
+
+    stack, nodata = build_rgbn(cube, mode_n, source)
+
+    return_probs = kwargs.get("return_probs", False)
+    pred = model.predict_array(
+        stack,
+        apply_nodata_mask=False,   # we apply the CHRIS nodata mask below
+        return_probs=return_probs,
+        verbose=kwargs.get("verbose", False),
+    )
+    if not return_probs:
+        pred = np.asarray(pred).astype(np.uint8)
+        pred[nodata] = 99
+    return pred