Redesign: five-step pedagogical flow with spectral baseline

Step 1 - Know Your Data: full-scene composite viewer (natural/false/SWIR/single-band)
with labelled-pixel markers overlaid, clickable pixel inspection.

Step 2 - Spectral Signatures: per-class mean±sigma chart and NDVI/NDWI heatmaps
computed from raw (unnormalised) bands for correct ratio semantics.

Step 3 - Spectral Baseline: pure-numpy chunked KNN classifier that predicts the
full 501x1001 scene from 7-band spectral features alone (no spatial context).

Step 4 - Deep Learning: UNet training with side-by-side ground-truth / KNN /
UNet patch comparison to make the spatial-context benefit concrete.

Step 5 - Experiment Lab: compare up to 3 UNet experiments; max capped at 3.

Also fixes natural-colour composite to use H4/H3/H2 (Red/Green/Blue) instead of
the previous H3/H2/H1, and adds matplotlib + scipy to requirements.

app.py

@@ -1,151 +1,270 @@
 import gradio as gr
 
-from config import APP_TITLE, set_seed, SEED, DEFAULT_PATCH_SIZE
+from config import (
+    APP_TITLE, set_seed, SEED, DEFAULT_PATCH_SIZE,
+    COMPOSITE_PRESETS, BAND_DESCRIPTIONS, MAX_EXPERIMENTS,
+)
 from train import (
     load_dataset_action,
-    update_explorer_sample,
-    update_compare_sample,
+    update_step1_composite,
+    handle_click_step1,
+    update_step2_index,
+    run_baseline_action,
+    update_step4_patch,
     train_experiment,
-    handle_click_dataset,
-    handle_click_exp_a,
-    handle_click_exp_b,
+    update_step5_comparison,
+    handle_click_step5,
 )
 
 set_seed(SEED)
 
+_COMPOSITE_CHOICES = list(COMPOSITE_PRESETS.keys()) + BAND_DESCRIPTIONS
+
 custom_css = """
-#compare-a img, #compare-b img, #explorer img {
-    image-rendering: pixelated;
-}
-.small-note { font-size: 0.9rem; opacity: 0.85; }
+#step1-img img, #step3-pred img, #step3-correct img { image-rendering: pixelated; }
+.step-header { font-size: 1.05rem; font-weight: 600; margin-bottom: 4px; }
+.hint { font-size: 0.88rem; color: #666; }
 """
 
 with gr.Blocks(title=APP_TITLE, css=custom_css) as demo:
-    gr.Markdown(f"# {APP_TITLE}\nInteractive teaching app for multispectral semantic segmentation.")
+    gr.Markdown(f"# {APP_TITLE}")
+    gr.Markdown(
+        "A five-step journey from raw satellite pixels to deep-learning segmentation. "
+        "Work through the tabs in order — each step builds on the previous one."
+    )
 
     dataset_state     = gr.State(None)
+    baseline_state    = gr.State(None)
     experiments_state = gr.State([])
 
-    # ── Tab 1: Image Explorer ────────────────────────────────
-    with gr.Tab("1) Image explorer"):
+    # ────────────────────────────────────────────────────────
+    # Step 1 — Know Your Data
+    # ────────────────────────────────────────────────────────
+    with gr.Tab("Step 1 · Know Your Data"):
+        gr.Markdown(
+            "**Start here.** Load the dataset, then explore the 7 spectral bands. "
+            "Squares on the image are training labels; circles are validation labels.",
+            elem_classes="hint",
+        )
         with gr.Row():
             with gr.Column(scale=1):
-                patch_size   = gr.Slider(64, 512, value=DEFAULT_PATCH_SIZE, step=32, label="Patch size")
+                patch_size   = gr.Slider(64, 512, value=DEFAULT_PATCH_SIZE, step=32, label="Patch size (for training)")
                 load_btn     = gr.Button("Load dataset", variant="primary")
-                dataset_info = gr.Markdown("### No dataset loaded yet")
+                composite_dd = gr.Dropdown(
+                    choices=_COMPOSITE_CHOICES,
+                    value="Natural Color (R/G/B)",
+                    label="View mode",
+                )
+                step1_info   = gr.Markdown("*Load the dataset to begin.*")
+
+            with gr.Column(scale=3, elem_id="step1-img"):
+                step1_image   = gr.Image(label="Full scene — click to inspect a pixel", type="numpy")
+                step1_click   = gr.Markdown("*Click anywhere on the image.*")
+
+    # ────────────────────────────────────────────────────────
+    # Step 2 — Spectral Signatures
+    # ────────────────────────────────────────────────────────
+    with gr.Tab("Step 2 · Spectral Signatures"):
+        gr.Markdown(
+            "Each land cover type has a characteristic pattern of brightness across the 7 bands — "
+            "its **spectral signature**. Notice how H_5 (NIR) separates vegetation from water. "
+            "NDVI and NDWI are hand-crafted indices that exploit this difference.",
+            elem_classes="hint",
+        )
+        with gr.Row():
+            with gr.Column(scale=1):
+                index_radio = gr.Radio(
+                    choices=["NDVI", "NDWI"],
+                    value="NDVI",
+                    label="Spectral index map",
+                )
                 gr.Markdown(
-                    "<div class='small-note'>Downloads the satellite dataset from HuggingFace "
-                    "and extracts image patches for training and validation.</div>"
+                    "**NDVI** = (NIR − Red) / (NIR + Red)  \n"
+                    "High values → dense vegetation (Forest, Agriculture)  \n"
+                    "Low / negative → water, urban, bare soil\n\n"
+                    "**NDWI** = (Green − NIR) / (Green + NIR)  \n"
+                    "Positive → water; negative → land",
+                    elem_classes="hint",
                 )
-            with gr.Column(scale=2, elem_id="explorer"):
-                explorer_sample_index = gr.Slider(0, 59, value=0, step=1, label="Validation patch index")
-                with gr.Row():
-                    explorer_rgb     = gr.Image(label="RGB composite",        type="numpy", height=400)
-                    explorer_gt      = gr.Image(label="Ground truth mask",    type="numpy", height=400)
-                    explorer_overlay = gr.Image(label="Ground truth overlay", type="numpy", height=400)
-                explorer_click_info = gr.Markdown("### Click the RGB image to inspect a pixel")
 
-    # ── Tab 2: Model Trainer ─────────────────────────────────
-    with gr.Tab("2) Model trainer"):
+            with gr.Column(scale=3):
+                step2_sig_chart  = gr.Image(label="Spectral signatures (training labels)", type="numpy")
+                step2_index_map  = gr.Image(label="Index map", type="numpy")
+
+    # ────────────────────────────────────────────────────────
+    # Step 3 — Spectral Baseline (KNN)
+    # ────────────────────────────────────────────────────────
+    with gr.Tab("Step 3 · Spectral Baseline"):
+        gr.Markdown(
+            "**No convolutions here.** We classify every pixel using only its 7 band values, "
+            "finding the k nearest training pixels in spectral space. "
+            "This shows you what's achievable without any spatial context.",
+            elem_classes="hint",
+        )
         with gr.Row():
             with gr.Column(scale=1):
-                run_name      = gr.Textbox(label="Experiment name", placeholder="e.g. lr-1e-3_ep-5")
-                learning_rate = gr.Slider(1e-4, 5e-3, value=1e-3, step=1e-4, label="Learning rate")
-                batch_size    = gr.Slider(2, 32,  value=8,  step=2,  label="Batch size")
-                epochs        = gr.Slider(1, 20,  value=5,  step=1,  label="Epochs")
-                base_channels = gr.Slider(8, 64,  value=16, step=8,  label="Model width (base channels)")
-                train_btn     = gr.Button("Train experiment", variant="primary")
+                k_slider        = gr.Slider(1, 5, value=3, step=2, label="k (number of neighbours)")
+                baseline_btn    = gr.Button("Run KNN baseline", variant="primary")
+                step3_metrics   = gr.Markdown("*Run the baseline to see results.*")
+
+            with gr.Column(scale=3, elem_id="step3-pred"):
+                step3_full_pred = gr.Image(
+                    label="Full-scene prediction · overlaid on natural-colour image · coloured dots = val labels (green=correct, red=wrong)",
+                    type="numpy",
+                )
+
+    # ────────────────────────────────────────────────────────
+    # Step 4 — Deep Learning (UNet)
+    # ────────────────────────────────────────────────────────
+    with gr.Tab("Step 4 · Deep Learning"):
+        gr.Markdown(
+            "A **U-Net** sees a patch of pixels at once, not just one pixel. "
+            "Its encoder captures local texture; skip connections preserve spatial detail. "
+            "Train a model and compare it patch-by-patch against the KNN baseline.",
+            elem_classes="hint",
+        )
+        with gr.Row():
             with gr.Column(scale=1):
-                train_summary = gr.Markdown("### No training run yet")
-                model_table   = gr.Markdown("### No models trained yet")
+                run_name      = gr.Textbox(label="Experiment name", placeholder="e.g. lr-1e-3_ch-16")
+                learning_rate = gr.Slider(1e-4, 5e-3, value=1e-3,  step=1e-4, label="Learning rate")
+                batch_size    = gr.Slider(2, 32,   value=8,  step=2, label="Batch size")
+                epochs        = gr.Slider(1, 20,   value=5,  step=1, label="Epochs")
+                base_channels = gr.Slider(8, 64,   value=16, step=8, label="Model width (base channels)")
+                train_btn     = gr.Button("Train model", variant="primary")
                 gr.Markdown(
-                    "<div class='small-note'>Accuracy and mIoU are computed on labeled pixels only "
-                    "(unlabeled pixels are ignored during training and evaluation).</div>"
+                    f"*Max {MAX_EXPERIMENTS} experiments total. "
+                    "Reload data to reset.*",
+                    elem_classes="hint",
                 )
+                step4_summary = gr.Markdown("*Train a model to see results.*")
 
-    # ── Tab 3: Result Comparison ─────────────────────────────
-    with gr.Tab("3) Result comparison"):
-        compare_sample_index = gr.Slider(0, 59, value=0, step=1, label="Validation patch index")
+            with gr.Column(scale=3):
+                step4_patch_slider = gr.Slider(0, 59, value=0, step=1, label="Validation patch index")
+                with gr.Row():
+                    step4_gt_img = gr.Image(label="Ground truth overlay",   type="numpy", height=280)
+                    step4_bl_img = gr.Image(label="KNN Baseline prediction", type="numpy", height=280)
+                    step4_un_img = gr.Image(label="UNet prediction",         type="numpy", height=280)
+
+    # ────────────────────────────────────────────────────────
+    # Step 5 — Experiment Lab
+    # ────────────────────────────────────────────────────────
+    with gr.Tab("Step 5 · Experiment Lab"):
+        gr.Markdown(
+            f"Compare up to **{MAX_EXPERIMENTS}** UNet experiments side by side. "
+            "Try different learning rates, epochs, or model widths and see what changes.",
+            elem_classes="hint",
+        )
         with gr.Row():
-            compare_sel_a = gr.Dropdown(choices=[], value=None, label="Left model",  interactive=True)
-            compare_sel_b = gr.Dropdown(choices=[], value=None, label="Right model", interactive=True)
+            step5_sel_a = gr.Dropdown(choices=[], value=None, label="Left model",  interactive=True)
+            step5_sel_b = gr.Dropdown(choices=[], value=None, label="Right model", interactive=True)
+        step5_patch_slider = gr.Slider(0, 59, value=0, step=1, label="Validation patch index")
+        step5_table = gr.Markdown("*No experiments yet.*")
+
+        gr.Markdown(
+            "**Guiding questions**\n"
+            "- Double the epochs — does mIoU keep improving or plateau?\n"
+            "- Halve the learning rate — does training become more stable?\n"
+            "- Increase base channels from 16 to 32 — worth the extra time?",
+            elem_classes="hint",
+        )
+
         with gr.Row():
-            with gr.Column(scale=1, elem_id="compare-a"):
+            with gr.Column(scale=1):
                 gr.Markdown("## Left")
-                compare_a_rgb     = gr.Image(label="RGB — click to inspect pixel", type="numpy", height=280)
-                compare_a_click   = gr.Markdown("### Click the RGB or overlay image to inspect a pixel")
-                compare_a_pred    = gr.Image(label="Prediction mask",    type="numpy", height=280)
-                compare_a_overlay = gr.Image(label="Prediction overlay — click to inspect pixel", type="numpy", height=280)
-                compare_a_metrics = gr.Markdown("### No model selected")
-                compare_a_error   = gr.Image(label="Correctness map",   type="numpy", height=280)
-
-            with gr.Column(scale=1, elem_id="compare-b"):
+                s5_a_rgb     = gr.Image(label="RGB",            type="numpy", height=240)
+                s5_a_pred    = gr.Image(label="Prediction",     type="numpy", height=240)
+                s5_a_overlay = gr.Image(label="Overlay",        type="numpy", height=240)
+                s5_a_metrics = gr.Markdown("*No model selected.*")
+                s5_a_error   = gr.Image(label="Correctness map",type="numpy", height=240)
+
+            with gr.Column(scale=1):
                 gr.Markdown("## Right")
-                compare_b_rgb     = gr.Image(label="RGB — click to inspect pixel", type="numpy", height=280)
-                compare_b_click   = gr.Markdown("### Click the RGB or overlay image to inspect a pixel")
-                compare_b_pred    = gr.Image(label="Prediction mask",    type="numpy", height=280)
-                compare_b_overlay = gr.Image(label="Prediction overlay — click to inspect pixel", type="numpy", height=280)
-                compare_b_metrics = gr.Markdown("### No model selected")
-                compare_b_error   = gr.Image(label="Correctness map",   type="numpy", height=280)
-
-    # ── Shared lists ─────────────────────────────────────────
-    # Order matches render_experiment_panel: (rgb, pred, overlay, metrics, error, click)
-    _compare_outputs = [
-        compare_a_rgb, compare_a_pred, compare_a_overlay, compare_a_metrics, compare_a_error, compare_a_click,
-        compare_b_rgb, compare_b_pred, compare_b_overlay, compare_b_metrics, compare_b_error, compare_b_click,
+                s5_b_rgb     = gr.Image(label="RGB",            type="numpy", height=240)
+                s5_b_pred    = gr.Image(label="Prediction",     type="numpy", height=240)
+                s5_b_overlay = gr.Image(label="Overlay",        type="numpy", height=240)
+                s5_b_metrics = gr.Markdown("*No model selected.*")
+                s5_b_error   = gr.Image(label="Correctness map",type="numpy", height=240)
+
+    # ── Event wiring ─────────────────────────────────────────
+
+    _load_outputs = [
+        dataset_state, baseline_state, experiments_state,
+        # Tab 1
+        step1_info, step1_image, step1_click,
+        # Tab 2
+        step2_sig_chart, step2_index_map,
+        # Tab 3
+        step3_metrics, step3_full_pred,
+        # Tab 4
+        step4_summary, step4_patch_slider,
+        step4_gt_img, step4_bl_img, step4_un_img,
+        # Tab 5
+        step5_table, step5_sel_a, step5_sel_b,
     ]
-    _compare_inputs = [dataset_state, experiments_state, compare_sel_a, compare_sel_b, compare_sample_index]
-
-    # ── Event bindings ────────────────────────────────────────
-
-    load_btn.click(
-        fn=load_dataset_action,
-        inputs=[patch_size],
-        outputs=[
-            dataset_state, experiments_state,
-            dataset_info,
-            explorer_rgb, explorer_gt, explorer_overlay,
-            explorer_click_info,
-            explorer_sample_index,
-            compare_sample_index,
-            compare_sel_a, compare_sel_b,
-        ],
+
+    load_btn.click(fn=load_dataset_action, inputs=[patch_size], outputs=_load_outputs)
+
+    composite_dd.change(
+        fn=update_step1_composite,
+        inputs=[dataset_state, composite_dd],
+        outputs=[step1_image, step1_click],
     )
 
-    explorer_sample_index.change(
-        fn=update_explorer_sample,
-        inputs=[dataset_state, explorer_sample_index],
-        outputs=[explorer_rgb, explorer_gt, explorer_overlay, explorer_click_info],
+    step1_image.select(
+        fn=handle_click_step1,
+        inputs=[dataset_state],
+        outputs=[step1_click],
     )
 
-    explorer_rgb.select(
-        fn=handle_click_dataset,
-        inputs=[dataset_state, explorer_sample_index],
-        outputs=[explorer_click_info],
+    index_radio.change(
+        fn=update_step2_index,
+        inputs=[dataset_state, index_radio],
+        outputs=[step2_index_map],
     )
 
+    baseline_btn.click(
+        fn=run_baseline_action,
+        inputs=[dataset_state, k_slider],
+        outputs=[baseline_state, step3_metrics, step3_full_pred],
+    )
+
+    _train_outputs = [
+        experiments_state,
+        step4_summary, step4_patch_slider,
+        step4_gt_img, step4_bl_img, step4_un_img,
+        step5_table, step5_sel_a, step5_sel_b,
+    ]
+
     train_btn.click(
         fn=train_experiment,
-        inputs=[dataset_state, experiments_state, learning_rate, batch_size, epochs, base_channels, run_name],
-        outputs=[experiments_state, train_summary, model_table, compare_sel_a, compare_sel_b],
+        inputs=[
+            dataset_state, baseline_state, experiments_state,
+            learning_rate, batch_size, epochs, base_channels, run_name,
+        ],
+        outputs=_train_outputs,
     )
 
-    for sel in [compare_sel_a, compare_sel_b]:
-        sel.change(fn=update_compare_sample, inputs=_compare_inputs, outputs=_compare_outputs)
+    step4_patch_slider.change(
+        fn=update_step4_patch,
+        inputs=[dataset_state, baseline_state, experiments_state, step4_patch_slider],
+        outputs=[step4_gt_img, step4_bl_img, step4_un_img],
+    )
 
-    compare_sample_index.change(fn=update_compare_sample, inputs=_compare_inputs, outputs=_compare_outputs)
+    _s5_inputs  = [dataset_state, experiments_state, step5_sel_a, step5_sel_b, step5_patch_slider]
+    _s5_outputs = [
+        s5_a_rgb, s5_a_pred, s5_a_overlay, s5_a_metrics, s5_a_error,
+        s5_b_rgb, s5_b_pred, s5_b_overlay, s5_b_metrics, s5_b_error,
+    ]
 
-    for img in [compare_a_rgb, compare_a_overlay]:
-        img.select(
-            fn=handle_click_exp_a,
-            inputs=[dataset_state, experiments_state, compare_sel_a, compare_sample_index],
-            outputs=[compare_a_click],
-        )
-    for img in [compare_b_rgb, compare_b_overlay]:
+    for trigger in [step5_sel_a, step5_sel_b, step5_patch_slider]:
+        trigger.change(fn=update_step5_comparison, inputs=_s5_inputs, outputs=_s5_outputs)
+
+    for img, sel in [(s5_a_rgb, step5_sel_a), (s5_a_overlay, step5_sel_a),
+                     (s5_b_rgb, step5_sel_b), (s5_b_overlay, step5_sel_b)]:
         img.select(
-            fn=handle_click_exp_b,
-            inputs=[dataset_state, experiments_state, compare_sel_b, compare_sample_index],
-            outputs=[compare_b_click],
+            fn=handle_click_step5,
+            inputs=[dataset_state, experiments_state, sel, step5_patch_slider],
+            outputs=[s5_a_metrics if sel == step5_sel_a else s5_b_metrics],
         )
 
 

baseline.py

@@ -0,0 +1,79 @@
+"""Spectral baseline classifier: KNN on raw 7-band pixel values, no spatial context."""
+import numpy as np
+
+from config import NUM_CHANNELS, NUM_CLASSES, IGNORE_INDEX
+from metrics import compute_metrics, metrics_markdown
+
+
+def _knn_predict(
+    train_X: np.ndarray,
+    train_y: np.ndarray,
+    query_X: np.ndarray,
+    k: int,
+    chunk: int = 50_000,
+) -> np.ndarray:
+    """Chunked nearest-neighbour prediction to keep peak RAM reasonable."""
+    N = len(query_X)
+    preds = np.empty(N, dtype=np.int64)
+    k = min(k, len(train_X))
+
+    for start in range(0, N, chunk):
+        end    = min(start + chunk, N)
+        block  = query_X[start:end]                                  # (B, 7)
+        dists  = np.sum((block[:, None, :] - train_X[None, :, :]) ** 2, axis=2)  # (B, N_tr)
+        nn_idx = np.argpartition(dists, k - 1, axis=1)[:, :k]       # (B, k)
+        labels = train_y[nn_idx]                                     # (B, k)
+        if k == 1:
+            preds[start:end] = labels[:, 0]
+        else:
+            # Vectorised majority vote
+            votes = (labels[:, :, None] == np.arange(NUM_CLASSES)[None, None, :]).sum(axis=1)
+            preds[start:end] = votes.argmax(axis=1)
+
+    return preds
+
+
+def run_knn_baseline(
+    full_image:      np.ndarray,
+    full_train_mask: np.ndarray,
+    full_val_mask:   np.ndarray,
+    val_images:      np.ndarray,
+    k: int = 3,
+):
+    """
+    Train KNN on labeled training pixels; predict (a) the full scene and (b) each
+    validation patch. Evaluate against the full validation mask.
+
+    Returns
+    -------
+    full_pred  : (H, W)         – class index for every pixel in the scene
+    val_preds  : (N, ph, pw)    – patch-level predictions for step-4 comparison
+    metrics    : dict
+    metrics_md : str
+    """
+    C, H, W = full_image.shape
+
+    labeled = full_train_mask != IGNORE_INDEX
+    if not labeled.any():
+        raise ValueError("No labeled training pixels found in TRAINING.tif.")
+
+    train_X = full_image[:, labeled].T            # (N_tr, 7)
+    train_y = full_train_mask[labeled]            # (N_tr,)
+
+    # --- Full scene prediction ---
+    all_X     = full_image.reshape(C, H * W).T   # (H*W, 7)
+    full_pred = _knn_predict(train_X, train_y, all_X, k).reshape(H, W)
+
+    # --- Validation patch predictions (same patches as UNet) ---
+    N_val, ph, pw = val_images.shape[0], val_images.shape[2], val_images.shape[3]
+    all_patch_X = np.concatenate(
+        [p.reshape(C, -1).T for p in val_images], axis=0
+    )  # (N_val * ph * pw, 7)
+    patch_preds_flat = _knn_predict(train_X, train_y, all_patch_X, k)
+    val_preds = patch_preds_flat.reshape(N_val, ph, pw).astype(np.int64)
+
+    # --- Metrics on full val mask ---
+    metrics    = compute_metrics(full_pred.ravel(), full_val_mask.ravel())
+    metrics_md = metrics_markdown(metrics, title=f"KNN Baseline  (k={k})")
+
+    return full_pred.astype(np.int64), val_preds, metrics, metrics_md

config.py

@@ -7,19 +7,36 @@ DEFAULT_PATCH_SIZE = 128
 NUM_CHANNELS = 7
 NUM_CLASSES = 4
 IGNORE_INDEX = 255
+MAX_EXPERIMENTS = 3
 
 BAND_NAMES = ["H_1", "H_2", "H_3", "H_4", "H_5", "H_6", "H_7"]
+BAND_DESCRIPTIONS = [
+    "H_1 (Coastal/Aerosol)",
+    "H_2 (Blue)",
+    "H_3 (Green)",
+    "H_4 (Red)",
+    "H_5 (NIR)",
+    "H_6 (SWIR-1)",
+    "H_7 (SWIR-2)",
+]
 CLASS_NAMES = ["Water", "Urban", "Agriculture", "Forest"]
 CLASS_COLORS = np.array(
     [
-        [30,  144, 255],   # Water       — blue
-        [220,  50,  50],   # Urban       — red
-        [255, 215,   0],   # Agriculture — yellow
-        [34,  139,  34],   # Forest      — green
+        [30,  144, 255],   # Water       - blue
+        [220,  50,  50],   # Urban       - red
+        [255, 215,   0],   # Agriculture - yellow
+        [34,  139,  34],   # Forest      - green
     ],
     dtype=np.uint8,
 )
 
+# (R-band-index, G-band-index, B-band-index) for composite presets
+COMPOSITE_PRESETS = {
+    "Natural Color (R/G/B)":          (3, 2, 1),
+    "False Color NIR (NIR/R/G)":      (4, 3, 2),
+    "SWIR Composite (SWIR2/NIR/R)":   (6, 4, 3),
+}
+
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 

data.py

@@ -19,9 +19,14 @@ VAL_MASK_FILE   = "GROUND TRUTH.tif"
 
 # ── File helpers ─────────────────────────────────────────────
 
-def _download(filename: str) -> str:
+def _get_path(filename: str) -> str:
+    """Use local file if it exists, otherwise download from HuggingFace."""
+    if os.path.exists(filename):
+        return filename
     if not DATASET_REPO:
-        raise EnvironmentError("DATASET_REPO not set in Space secrets.")
+        raise EnvironmentError(
+            f"'{filename}' not found locally and DATASET_REPO is not set."
+        )
     return hf_hub_download(
         repo_id=DATASET_REPO,
         filename=filename,
@@ -39,11 +44,10 @@ def _read_band(path: str) -> np.ndarray:
 
 
 def _read_mask_raw(path: str) -> Tuple[np.ndarray, object, str]:
-    """Returns (raw_array, nodata_value, info_string)."""
     with rasterio.open(path) as src:
         data   = src.read(1)
         nodata = src.nodata
-        info   = f"shape={src.shape} dtype={src.dtypes[0]} nodata={nodata} bands={src.count}"
+        info   = f"shape={src.shape} dtype={src.dtypes[0]} nodata={nodata}"
     return data, nodata, info
 
 
@@ -59,22 +63,33 @@ def _normalize(image: np.ndarray) -> np.ndarray:
     return out
 
 
+def _compute_ndvi(raw_image: np.ndarray) -> np.ndarray:
+    """NDVI = (NIR - Red) / (NIR + Red).  H_5=NIR (idx 4), H_4=Red (idx 3)."""
+    nir = raw_image[4].astype(np.float32)
+    red = raw_image[3].astype(np.float32)
+    denom = nir + red
+    return np.where(np.abs(denom) > 1e-6, (nir - red) / denom, 0.0).clip(-1.0, 1.0)
+
+
+def _compute_ndwi(raw_image: np.ndarray) -> np.ndarray:
+    """NDWI = (Green - NIR) / (Green + NIR).  H_3=Green (idx 2), H_5=NIR (idx 4)."""
+    green = raw_image[2].astype(np.float32)
+    nir   = raw_image[4].astype(np.float32)
+    denom = green + nir
+    return np.where(np.abs(denom) > 1e-6, (green - nir) / denom, 0.0).clip(-1.0, 1.0)
+
+
 def _remap_mask(raw: np.ndarray, nodata_val) -> Tuple[np.ndarray, List[int]]:
     """
-    Map raw pixel values to 0..NUM_CLASSES-1.
-    Value 0 is treated as unlabeled background → IGNORE_INDEX.
-    Nodata pixels → IGNORE_INDEX.
-    Returns (remapped_mask, sorted_raw_class_values_used).
+    Map raw pixel values -> 0..NUM_CLASSES-1.
+    Value 0 and nodata -> IGNORE_INDEX.
     """
     if nodata_val is not None:
         nodata_px = raw == int(nodata_val)
     else:
         nodata_px = np.zeros(raw.shape, dtype=bool)
 
-    # Treat pixel value 0 as unlabeled background
-    background_px = raw == 0
-    ignore_px = nodata_px | background_px
-
+    ignore_px  = nodata_px | (raw == 0)
     valid      = ~ignore_px
     raw_unique = sorted(int(v) for v in np.unique(raw[valid]))
 
@@ -97,7 +112,6 @@ def _extract_patches(
     stride = patch_size // 2
     imgs, masks = [], []
 
-    # Build step lists that always include the last valid position (covers edges)
     def steps(size):
         s = list(range(0, size - patch_size + 1, stride))
         if not s:
@@ -110,12 +124,10 @@ def _extract_patches(
         for x in steps(W):
             pm = mask[y : y + patch_size, x : x + patch_size]
             pi = image[:, y : y + patch_size, x : x + patch_size]
-            # Include any patch that contains at least one labeled pixel
             if pm.shape == (patch_size, patch_size) and (pm != IGNORE_INDEX).any():
                 imgs.append(pi)
                 masks.append(pm)
 
-    # Last resort: pad with zeros/IGNORE if image is smaller than patch_size
     if not imgs:
         ph = min(patch_size, H)
         pw = min(patch_size, W)
@@ -129,6 +141,29 @@ def _extract_patches(
     return np.stack(imgs).astype(np.float32), np.stack(masks).astype(np.int64)
 
 
+# ── Spectral analysis ─────────────────────────────────────────
+
+def compute_spectral_signatures(full_image: np.ndarray, full_mask: np.ndarray) -> Dict:
+    """Per-class mean and std across the 7 normalized bands, from labeled pixels."""
+    sigs = {}
+    for cls_idx in range(NUM_CLASSES):
+        px = full_mask == cls_idx
+        if px.sum() == 0:
+            sigs[cls_idx] = {
+                "mean": np.zeros(NUM_CHANNELS, dtype=np.float32),
+                "std":  np.zeros(NUM_CHANNELS, dtype=np.float32),
+                "n":    0,
+            }
+        else:
+            vals = full_image[:, px]  # (7, N)
+            sigs[cls_idx] = {
+                "mean": vals.mean(axis=1).astype(np.float32),
+                "std":  vals.std(axis=1).astype(np.float32),
+                "n":    int(px.sum()),
+            }
+    return sigs
+
+
 # ── Dataset class ─────────────────────────────────────────────
 
 class MultiSpectralDataset(Dataset):
@@ -146,28 +181,38 @@ class MultiSpectralDataset(Dataset):
 # ── Public API ────────────────────────────���───────────────────
 
 def load_data(patch_size: int = DEFAULT_PATCH_SIZE) -> Dict:
-    # Download and stack bands
-    band_arrays = [_read_band(_download(f)) for f in BAND_FILES]
-    image = _normalize(np.stack(band_arrays, axis=0))  # (7, H, W) float32
+    # Read raw bands
+    raw_bands = [_read_band(_get_path(f)) for f in BAND_FILES]
+    raw_image = np.stack(raw_bands, axis=0)  # (7, H, W) raw float32
+
+    # Compute spectral indices from raw values (ratio, so must use raw)
+    ndvi = _compute_ndvi(raw_image)
+    ndwi = _compute_ndwi(raw_image)
 
-    # Read raw masks + metadata
-    raw_train, nd_train, info_train = _read_mask_raw(_download(TRAIN_MASK_FILE))
-    raw_val,   nd_val,   info_val   = _read_mask_raw(_download(VAL_MASK_FILE))
+    # Normalize
+    image = _normalize(raw_image)  # (7, H, W) normalized [0,1]
 
-    # Remap to 0-indexed classes
+    # Read masks
+    raw_train, nd_train, info_train = _read_mask_raw(_get_path(TRAIN_MASK_FILE))
+    raw_val,   nd_val,   info_val   = _read_mask_raw(_get_path(VAL_MASK_FILE))
+
+    # Remap
     train_mask, train_vals = _remap_mask(raw_train, nd_train)
     val_mask,   val_vals   = _remap_mask(raw_val,   nd_val)
 
     if not train_vals:
         raise ValueError(
-            f"TRAINING.tif has no labeled pixels after nodata removal. "
-            f"File info: {info_train} | Unique raw values: {np.unique(raw_train).tolist()}"
+            f"TRAINING.tif has no labeled pixels. Info: {info_train} | "
+            f"Unique raw values: {np.unique(raw_train).tolist()}"
         )
 
-    # Extract patches
+    # Patches
     tr_imgs, tr_masks = _extract_patches(image, train_mask, patch_size)
     va_imgs, va_masks = _extract_patches(image, val_mask,   patch_size)
 
+    # Spectral signatures from training labels
+    signatures = compute_spectral_signatures(image, train_mask)
+
     train_labeled = int((train_mask != IGNORE_INDEX).sum())
     val_labeled   = int((val_mask   != IGNORE_INDEX).sum())
 
@@ -179,21 +224,21 @@ def load_data(patch_size: int = DEFAULT_PATCH_SIZE) -> Dict:
         return " | ".join(parts)
 
     status = "\n".join([
-        f"Train patches: **{len(tr_imgs)}** | Val patches: **{len(va_imgs)}** | Patch: **{patch_size}×{patch_size}**",
-        "",
-        f"**TRAINING.tif** `{info_train}`",
-        f"Raw values → classes: `{dict(zip(train_vals, CLASS_NAMES[:len(train_vals)]))}`",
-        f"Labeled pixels: **{train_labeled:,}** — {_class_dist(train_mask, train_labeled)}",
-        "",
-        f"**GROUND TRUTH.tif** `{info_val}`",
-        f"Raw values → classes: `{dict(zip(val_vals, CLASS_NAMES[:len(val_vals)]))}`",
-        f"Labeled pixels: **{val_labeled:,}** — {_class_dist(val_mask, val_labeled)}",
+        f"Train patches: **{len(tr_imgs)}** | Val patches: **{len(va_imgs)}** | Patch: **{patch_size}x{patch_size}**",
+        f"Training labels: **{train_labeled:,}** px — {_class_dist(train_mask, train_labeled)}",
+        f"Validation labels: **{val_labeled:,}** px — {_class_dist(val_mask, val_labeled)}",
     ])
 
     return {
-        "train_images": tr_imgs,
-        "train_masks":  tr_masks,
-        "val_images":   va_imgs,
-        "val_masks":    va_masks,
-        "status":       status,
+        "full_image":       image,
+        "full_train_mask":  train_mask,
+        "full_val_mask":    val_mask,
+        "ndvi":             ndvi,
+        "ndwi":             ndwi,
+        "signatures":       signatures,
+        "train_images":     tr_imgs,
+        "train_masks":      tr_masks,
+        "val_images":       va_imgs,
+        "val_masks":        va_masks,
+        "status":           status,
     }

requirements.txt

@@ -4,3 +4,5 @@ Pillow
 torch
 rasterio
 huggingface_hub
+matplotlib
+scipy

train.py

@@ -9,11 +9,20 @@ import gradio as gr
 
 from config import (
     DEVICE, NUM_CHANNELS, NUM_CLASSES, DEFAULT_PATCH_SIZE,
-    BAND_NAMES, CLASS_NAMES, IGNORE_INDEX,
+    BAND_NAMES, BAND_DESCRIPTIONS, CLASS_NAMES, IGNORE_INDEX,
+    COMPOSITE_PRESETS, MAX_EXPERIMENTS,
 )
 from data import MultiSpectralDataset, load_data
 from model import SmallUNet
-from visualize import multispectral_to_rgb, mask_to_color, overlay_mask, correctness_overlay
+from baseline import run_knn_baseline
+from visualize import (
+    render_composite, render_single_band,
+    add_labels_overlay, multispectral_to_rgb,
+    mask_to_color, overlay_mask, correctness_overlay,
+    render_full_prediction_overlay,
+    render_spectral_signatures_chart, render_index_map,
+    _blank_rgb,
+)
 from metrics import compute_metrics, metrics_markdown
 
 
@@ -22,8 +31,8 @@ from metrics import compute_metrics, metrics_markdown
 def build_prediction_cache(
     model: nn.Module, images: np.ndarray, batch_size: int = 8
 ) -> Tuple[np.ndarray, np.ndarray]:
-    dummy_masks = np.zeros((len(images), images.shape[-2], images.shape[-1]), dtype=np.int64)
-    ds     = MultiSpectralDataset(images, dummy_masks)
+    dummy = np.zeros((len(images), images.shape[-2], images.shape[-1]), dtype=np.int64)
+    ds     = MultiSpectralDataset(images, dummy)
     loader = DataLoader(ds, batch_size=batch_size, shuffle=False)
     preds, probs = [], []
     model.eval()
@@ -36,16 +45,19 @@ def build_prediction_cache(
     return np.concatenate(preds, axis=0), np.concatenate(probs, axis=0)
 
 
-# ── Render helpers ───────────────────────────────────────────
+# ── Shared render helpers ────────────────────────────────────
 
-def _blank(size: int = DEFAULT_PATCH_SIZE) -> np.ndarray:
-    return np.full((size, size, 3), 200, dtype=np.uint8)
+def _get_exp_by_name(experiments: List[Dict], name: Optional[str]) -> Optional[Dict]:
+    if not name:
+        return None
+    return next((e for e in experiments if e["name"] == name), None)
 
 
 def pixel_info_markdown(
     x: int, y: int,
     img7: np.ndarray, gt: np.ndarray,
-    pred: Optional[np.ndarray], probs: Optional[np.ndarray],
+    pred: Optional[np.ndarray] = None,
+    probs: Optional[np.ndarray] = None,
 ) -> str:
     h, w = gt.shape
     x = int(np.clip(x, 0, w - 1))
@@ -53,165 +65,244 @@ def pixel_info_markdown(
 
     gt_class = int(gt[y, x])
     gt_name  = CLASS_NAMES[gt_class] if gt_class != IGNORE_INDEX else "Unlabeled"
-    lines = [f"### Pixel ({x}, {y})", f"- Ground truth: **{gt_name}**"]
+    lines    = [f"**Pixel ({x}, {y})**", f"Ground truth: **{gt_name}**"]
 
     if pred is not None:
+        pred_class = int(pred[y, x])
+        lines.append(f"Prediction: **{CLASS_NAMES[pred_class]}**")
         if gt_class != IGNORE_INDEX:
-            pred_class = int(pred[y, x])
-            lines.append(f"- Prediction: **{CLASS_NAMES[pred_class]}**")
-            lines.append(f"- Correct: **{'Yes' if pred_class == gt_class else 'No'}**")
-            if probs is not None:
-                top_ids = np.argsort(probs[:, y, x])[::-1][:3]
-                lines.append("- Top probabilities: " + ", ".join(
-                    f"{CLASS_NAMES[i]} {probs[i, y, x] * 100:.1f}%" for i in top_ids
-                ))
-        else:
-            lines.append("- Prediction: —  *(unlabeled pixel)*")
-    else:
-        lines.append("- Prediction: —")
+            lines.append(f"Correct: **{'Yes' if pred_class == gt_class else 'No'}**")
+        if probs is not None:
+            top = np.argsort(probs[:, y, x])[::-1][:3]
+            lines.append("Top probs: " + ", ".join(
+                f"{CLASS_NAMES[i]} {probs[i, y, x]*100:.1f}%" for i in top
+            ))
 
     lines += ["", "**Band values**"] + [
-        f"- {n}: {float(img7[b, y, x]):.3f}" for b, n in enumerate(BAND_NAMES)
+        f"{BAND_DESCRIPTIONS[b]}: **{float(img7[b, y, x]):.3f}**"
+        for b in range(img7.shape[0])
+    ]
+    return "\n\n".join(lines)
+
+
+def experiments_table_markdown(experiments: List[Dict]) -> str:
+    if not experiments:
+        return "No experiments trained yet."
+    lines = [
+        "| # | Name | LR | Epochs | Channels | Val Acc | mIoU |",
+        "|---|---|---:|---:|---:|---:|---:|",
     ]
+    for i, e in enumerate(experiments):
+        cfg = e["config"]
+        lines.append(
+            f"| {i+1} | {e['name']} | {cfg['learning_rate']:.4f} | {cfg['epochs']} "
+            f"| {cfg['base_channels']} "
+            f"| {e['global_metrics']['overall_acc']*100:.1f}% "
+            f"| {e['global_metrics']['miou']*100:.1f}% |"
+        )
     return "\n".join(lines)
 
 
-def _get_exp_by_name(experiments: List[Dict], name: Optional[str]) -> Optional[Dict]:
-    if not name:
-        return None
-    return next((e for e in experiments if e["name"] == name), None)
+# ── Step 1 render helpers ──────────────────────────��─────────
+
+def _render_step1_image(dataset_state: Dict, composite_choice: str) -> np.ndarray:
+    full = dataset_state["full_image"]
+    if composite_choice in COMPOSITE_PRESETS:
+        r, g, b = COMPOSITE_PRESETS[composite_choice]
+        base = render_composite(full, r, g, b)
+    else:
+        band_idx = BAND_DESCRIPTIONS.index(composite_choice)
+        base = render_single_band(full, band_idx)
+    return add_labels_overlay(base, dataset_state["full_train_mask"], dataset_state["full_val_mask"])
+
 
+# ── Step 4 render helpers ────────────────────────────────────
 
-def render_experiment_panel(
-    dataset_state: Dict, exp: Optional[Dict], sample_idx: int
-) -> Tuple:
-    """Returns (rgb, pred_color, overlay, metrics_md, error_map, click_md)."""
-    b      = _blank()
-    no_data = (b, b, b, "### No data loaded", b, "### Load a dataset first")
-    if dataset_state is None or "val_images" not in dataset_state:
-        return no_data
+def _render_step4_row(
+    dataset_state: Dict,
+    baseline_state: Optional[Dict],
+    experiments: List[Dict],
+    patch_idx: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Returns (rgb, gt_overlay, baseline_overlay, unet_overlay)."""
     val_images = dataset_state["val_images"]
     val_masks  = dataset_state["val_masks"]
-    if len(val_images) == 0:
-        return no_data
+    idx = max(0, min(patch_idx, len(val_images) - 1))
 
-    idx = max(0, min(int(sample_idx), len(val_images) - 1))
     rgb = multispectral_to_rgb(val_images[idx])
     gt  = val_masks[idx]
+    gt_ov = overlay_mask(rgb, gt)
 
-    if exp is None:
-        return (
-            rgb, mask_to_color(gt), overlay_mask(rgb, gt),
-            "### No model selected",
-            _blank(),
-            pixel_info_markdown(0, 0, val_images[idx], gt, None, None),
-        )
+    if baseline_state is not None and idx < len(baseline_state["val_preds"]):
+        bl_ov = overlay_mask(rgb, baseline_state["val_preds"][idx])
+    else:
+        bl_ov = _blank_rgb(*rgb.shape[:2])
+
+    if experiments and idx < len(experiments[-1]["val_preds"]):
+        un_ov = overlay_mask(rgb, experiments[-1]["val_preds"][idx])
+    else:
+        un_ov = _blank_rgb(*rgb.shape[:2])
+
+    return rgb, gt_ov, bl_ov, un_ov
+
+
+# ── Step 5 render helpers ────────────────────────────────────
+
+def render_step5_panel(
+    dataset_state: Dict,
+    exp: Optional[Dict],
+    patch_idx: int,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, str, np.ndarray]:
+    """Returns (rgb, pred_color, overlay, metrics_md, error_map)."""
+    blank = _blank_rgb()
+    if dataset_state is None or exp is None:
+        return blank, blank, blank, "No model selected.", blank
+
+    val_images = dataset_state["val_images"]
+    val_masks  = dataset_state["val_masks"]
+    idx = max(0, min(patch_idx, len(val_images) - 1))
+    rgb = multispectral_to_rgb(val_images[idx])
+    gt  = val_masks[idx]
 
     if idx >= len(exp["val_preds"]):
-        return (
-            rgb, mask_to_color(gt), overlay_mask(rgb, gt),
-            "### Dataset reloaded — retrain to refresh",
-            _blank(),
-            "### Retrain needed",
-        )
+        return rgb, mask_to_color(gt), overlay_mask(rgb, gt), "Dataset reloaded — retrain.", blank
 
     pred  = exp["val_preds"][idx].astype(np.int64)
     probs = exp["val_probs"][idx].astype(np.float32)
-    sample_metrics = compute_metrics(pred, gt, num_classes=NUM_CLASSES)
+    m     = compute_metrics(pred, gt)
     return (
         rgb,
         mask_to_color(pred),
         overlay_mask(rgb, pred),
-        metrics_markdown(sample_metrics, title=f"{exp['name']} (sample {idx})"),
+        metrics_markdown(m, title=f"{exp['name']} · patch {idx}"),
         correctness_overlay(rgb, pred, gt),
-        pixel_info_markdown(0, 0, val_images[idx], gt, pred, probs),
-    )
-
-
-def render_compare_view(
-    dataset_state, experiments, name_a, name_b, sample_idx: int
-) -> Tuple:
-    return (
-        *render_experiment_panel(dataset_state, _get_exp_by_name(experiments, name_a), sample_idx),
-        *render_experiment_panel(dataset_state, _get_exp_by_name(experiments, name_b), sample_idx),
     )
 
 
-def experiments_table_markdown(experiments: List[Dict]) -> str:
-    if not experiments:
-        return "### No models trained yet"
-    lines = [
-        "### Trained models", "",
-        "| # | Name | LR | Epochs | Base Ch | Val Acc | mIoU |",
-        "|---|---|---:|---:|---:|---:|---:|",
-    ]
-    for i, e in enumerate(experiments):
-        cfg = e["config"]
-        lines.append(
-            f"| {i + 1} | {e['name']} | {cfg['learning_rate']:.4f} | {cfg['epochs']} "
-            f"| {cfg['base_channels']} | {e['global_metrics']['overall_acc'] * 100:.1f}% "
-            f"| {e['global_metrics']['miou'] * 100:.1f}% |"
-        )
-    return "\n".join(lines)
-
-
 # ── Gradio action functions ─────���────────────────────────────
 
 def load_dataset_action(patch_size: int):
-    """Downloads the HF dataset and builds patches. Returns 9 values."""
-    patch_size   = int(patch_size)
+    patch_size    = int(patch_size)
     dataset_state = load_data(patch_size)
-    val_count    = len(dataset_state["val_images"])
-
-    rgb = multispectral_to_rgb(dataset_state["val_images"][0])
-    gt  = dataset_state["val_masks"][0]
-    dataset_info = "\n".join([
-        "### Dataset loaded",
-        f"- {dataset_state['status']}",
-        f"- Channels: **{NUM_CHANNELS}** ({', '.join(BAND_NAMES)})",
-        f"- Classes: **{NUM_CLASSES}** ({', '.join(CLASS_NAMES)})",
+    val_count     = len(dataset_state["val_images"])
+
+    step1_img   = _render_step1_image(dataset_state, "Natural Color (R/G/B)")
+    sig_chart   = render_spectral_signatures_chart(dataset_state["signatures"])
+    ndvi_map    = render_index_map(
+        dataset_state["ndvi"], "NDVI",
+        dataset_state["full_train_mask"], dataset_state["full_val_mask"],
+    )
+    blank = _blank_rgb()
+
+    dataset_info = "\n\n".join([
+        "**Dataset loaded.**",
+        dataset_state["status"],
+        f"Bands: {', '.join(BAND_NAMES)}  |  Classes: {', '.join(CLASS_NAMES)}",
+        "**Squares** = training labels · **Circles** = validation labels",
     ])
 
+    slider_upd = gr.update(maximum=max(0, val_count - 1), value=0)
+
     return (
         dataset_state,
-        [],
+        None,                               # baseline_state reset
+        [],                                 # experiments_state reset
+        # Tab 1
         dataset_info,
-        rgb,
-        mask_to_color(gt),
-        overlay_mask(rgb, gt),
-        pixel_info_markdown(0, 0, dataset_state["val_images"][0], gt, None, None),
-        gr.update(maximum=max(0, val_count - 1), value=0),   # explorer_sample_index
-        gr.update(maximum=max(0, val_count - 1), value=0),   # compare_sample_index
-        gr.update(choices=[], value=None),                    # compare_sel_a
-        gr.update(choices=[], value=None),                    # compare_sel_b
+        step1_img,
+        "Click the image to inspect a pixel.",
+        # Tab 2
+        sig_chart,
+        ndvi_map,
+        # Tab 3
+        "Run KNN baseline after loading the dataset.",
+        blank,
+        # Tab 4
+        "Train a model in Step 4.",
+        slider_upd,
+        blank, blank, blank,
+        # Tab 5
+        "No experiments yet.",
+        gr.update(choices=[], value=None),
+        gr.update(choices=[], value=None),
     )
 
 
-def update_explorer_sample(dataset_state, sample_idx: int):
-    if dataset_state is None or "val_images" not in dataset_state:
-        b = _blank()
-        return b, b, b, "### No dataset loaded"
-    val_images = dataset_state["val_images"]
-    val_masks  = dataset_state["val_masks"]
-    idx = max(0, min(int(sample_idx), len(val_images) - 1))
-    rgb = multispectral_to_rgb(val_images[idx])
-    gt  = val_masks[idx]
-    return (
-        rgb,
-        mask_to_color(gt),
-        overlay_mask(rgb, gt),
-        pixel_info_markdown(0, 0, val_images[idx], gt, None, None),
+def update_step1_composite(dataset_state, composite_choice: str):
+    if dataset_state is None:
+        return _blank_rgb(), "Load the dataset first."
+    img = _render_step1_image(dataset_state, composite_choice)
+    return img, "Click the image to inspect a pixel."
+
+
+def handle_click_step1(evt: gr.SelectData, dataset_state):
+    if dataset_state is None:
+        return "Load the dataset first."
+    x, y = evt.index
+    full  = dataset_state["full_image"]
+    fmask = dataset_state["full_val_mask"]
+    H, W  = fmask.shape
+    x, y  = int(np.clip(x, 0, W-1)), int(np.clip(y, 0, H-1))
+
+    cls   = int(fmask[y, x])
+    label = CLASS_NAMES[cls] if cls != IGNORE_INDEX else "Unlabeled"
+    lines = [
+        f"**Pixel ({x}, {y})**  |  Val label: **{label}**", "",
+        "| Band | Value |", "|---|---:|",
+    ] + [f"| {BAND_DESCRIPTIONS[b]} | {float(full[b, y, x]):.4f} |" for b in range(7)]
+    return "\n".join(lines)
+
+
+def update_step2_index(dataset_state, index_choice: str):
+    if dataset_state is None:
+        return _blank_rgb()
+    key = index_choice.lower()
+    arr = dataset_state[key]
+    return render_index_map(
+        arr, index_choice,
+        dataset_state["full_train_mask"], dataset_state["full_val_mask"],
     )
 
 
-def update_compare_sample(dataset_state, experiments, sel_a, sel_b, sample_idx: int):
-    if dataset_state is None or "val_images" not in dataset_state:
-        raise gr.Error("Load a dataset first.")
-    return render_compare_view(dataset_state, experiments, sel_a, sel_b, int(sample_idx))
+def run_baseline_action(dataset_state, k: int, progress=gr.Progress()):
+    if dataset_state is None:
+        raise gr.Error("Load the dataset first.")
+    progress(0.1, desc="Running KNN on full scene...")
+    k = int(k)
+    full_pred, val_preds, metrics, metrics_md = run_knn_baseline(
+        dataset_state["full_image"],
+        dataset_state["full_train_mask"],
+        dataset_state["full_val_mask"],
+        dataset_state["val_images"],
+        k=k,
+    )
+    progress(0.9, desc="Rendering...")
+    baseline_state = {
+        "k":         k,
+        "full_pred": full_pred,
+        "val_preds": val_preds,
+        "metrics":   metrics,
+    }
+    full_ov = render_full_prediction_overlay(
+        dataset_state["full_image"], full_pred, dataset_state["full_val_mask"],
+    )
+    progress(1.0)
+    return baseline_state, metrics_md, full_ov
+
+
+def update_step4_patch(dataset_state, baseline_state, experiments, patch_idx: int):
+    if dataset_state is None:
+        blank = _blank_rgb()
+        return blank, blank, blank
+    _, gt_ov, bl_ov, un_ov = _render_step4_row(
+        dataset_state, baseline_state, experiments, int(patch_idx)
+    )
+    return gt_ov, bl_ov, un_ov
 
 
 def train_experiment(
     dataset_state: Dict,
+    baseline_state: Optional[Dict],
     experiments: List[Dict],
     learning_rate: float,
     batch_size: int,
@@ -222,6 +313,11 @@ def train_experiment(
 ):
     if dataset_state is None or "train_images" not in dataset_state:
         raise gr.Error("Load a dataset first.")
+    if len(experiments) >= MAX_EXPERIMENTS:
+        raise gr.Error(
+            f"Maximum {MAX_EXPERIMENTS} experiments reached. "
+            "Go to Step 5 to compare, then reload data to start fresh."
+        )
 
     loader = DataLoader(
         MultiSpectralDataset(dataset_state["train_images"], dataset_state["train_masks"]),
@@ -233,39 +329,39 @@ def train_experiment(
 
     n_epochs = int(epochs)
     history  = []
-    for epoch_i in range(n_epochs):
-        progress(epoch_i / n_epochs, desc=f"Epoch {epoch_i + 1}/{n_epochs}")
+    for ep in range(n_epochs):
+        progress(ep / n_epochs, desc=f"Epoch {ep+1}/{n_epochs}")
         model.train()
-        total_loss, n = 0.0, 0
+        total, n = 0.0, 0
         for xb, yb in loader:
             xb, yb = xb.to(DEVICE), yb.to(DEVICE)
             optimizer.zero_grad(set_to_none=True)
             loss = criterion(model(xb), yb)
             loss.backward()
             optimizer.step()
-            total_loss += float(loss.item())
+            total += float(loss.item())
             n += 1
-        history.append(total_loss / max(1, n))
+        history.append(total / max(1, n))
 
     progress(0.95, desc="Running validation inference...")
     val_preds, val_probs = build_prediction_cache(
         model, dataset_state["val_images"], batch_size=max(1, int(batch_size))
     )
     global_metrics = compute_metrics(
-        val_preds.reshape(-1), dataset_state["val_masks"].reshape(-1), num_classes=NUM_CLASSES
+        val_preds.reshape(-1), dataset_state["val_masks"].reshape(-1)
     )
-    progress(1.0, desc="Done!")
+    progress(1.0)
 
-    base = (run_name or f"Run {len(experiments) + 1}").strip()
+    base = (run_name or f"Run {len(experiments)+1}").strip()
     existing = {e["name"] for e in experiments}
-    name, counter = base, 2
+    name, ctr = base, 2
     while name in existing:
-        name = f"{base} ({counter})"
-        counter += 1
+        name = f"{base} ({ctr})"
+        ctr += 1
 
     experiment = {
-        "name": name,
-        "config": {
+        "name":               name,
+        "config":             {
             "learning_rate": float(learning_rate),
             "batch_size":    int(batch_size),
             "epochs":        int(epochs),
@@ -273,63 +369,59 @@ def train_experiment(
         },
         "train_loss_history": history,
         "global_metrics":     global_metrics,
-        "val_preds":  val_preds.astype(np.int64),
-        "val_probs":  val_probs.astype(np.float32),
+        "val_preds":          val_preds.astype(np.int64),
+        "val_probs":          val_probs.astype(np.float32),
     }
-
     experiments = experiments + [experiment]
-    summary = "\n".join([
-        f"### Training finished — **{name}**",
-        f"- Device: **{DEVICE}** | Epochs: **{n_epochs}**",
-        f"- Final loss: **{history[-1]:.4f}**",
-        f"- Val accuracy: **{global_metrics['overall_acc'] * 100:.2f}%** (labeled pixels only)",
-        f"- Val mIoU: **{global_metrics['miou'] * 100:.2f}%**",
+
+    summary = "\n\n".join([
+        f"**Training finished — {name}**",
+        f"Device: **{DEVICE}** | Epochs: **{n_epochs}** | Final loss: **{history[-1]:.4f}**",
+        f"Val accuracy: **{global_metrics['overall_acc']*100:.2f}%** (labeled px only)",
+        f"Val mIoU: **{global_metrics['miou']*100:.2f}%**",
     ])
 
-    choices = [e["name"] for e in experiments]
+    choices   = [e["name"] for e in experiments]
+    val_count = len(dataset_state["val_images"])
+
+    _, gt_ov, bl_ov, un_ov = _render_step4_row(dataset_state, baseline_state, experiments, 0)
+
     return (
-        experiments, summary,
-        experiments_table_markdown(experiments),
-        gr.update(choices=choices),
-        gr.update(choices=choices),
+        experiments,
+        summary,
+        gr.update(maximum=max(0, val_count-1), value=0),   # step4 patch slider
+        gt_ov, bl_ov, un_ov,
+        experiments_table_markdown(experiments),            # step5 table
+        gr.update(choices=choices, value=None),             # step5 sel_a
+        gr.update(choices=choices, value=None),             # step5 sel_b
     )
 
 
-# ── Click handlers ───────────────────────────────────────────
-
-def handle_click_dataset(evt: gr.SelectData, dataset_state, sample_idx: int):
-    if dataset_state is None or "val_images" not in dataset_state:
-        return "### No dataset"
-    idx = max(0, min(int(sample_idx), len(dataset_state["val_images"]) - 1))
-    x, y = evt.index
-    return pixel_info_markdown(
-        int(x), int(y),
-        dataset_state["val_images"][idx], dataset_state["val_masks"][idx],
-        None, None,
-    )
+def update_step5_comparison(
+    dataset_state, experiments, sel_a, sel_b, patch_idx: int
+):
+    idx  = int(patch_idx)
+    exp_a = _get_exp_by_name(experiments, sel_a)
+    exp_b = _get_exp_by_name(experiments, sel_b)
+    a_outs = render_step5_panel(dataset_state, exp_a, idx)
+    b_outs = render_step5_panel(dataset_state, exp_b, idx)
+    return (*a_outs, *b_outs)
 
 
-def _handle_click_experiment(
-    evt: gr.SelectData, dataset_state, experiments,
-    model_name: Optional[str], sample_idx: int,
+def handle_click_step5(
+    evt: gr.SelectData,
+    dataset_state, experiments, model_name, patch_idx: int,
 ) -> str:
     try:
-        if dataset_state is None or "val_images" not in dataset_state:
-            return "### No dataset loaded"
-        idx = max(0, min(int(sample_idx), len(dataset_state["val_images"]) - 1))
+        if dataset_state is None:
+            return "No dataset loaded."
+        idx = max(0, min(int(patch_idx), len(dataset_state["val_images"])-1))
         exp = _get_exp_by_name(experiments, model_name)
         x, y = evt.index
         img7 = dataset_state["val_images"][idx]
         gt   = dataset_state["val_masks"][idx]
-        if exp is None or idx >= len(exp["val_preds"]):
-            return pixel_info_markdown(int(x), int(y), img7, gt, None, None)
-        return pixel_info_markdown(int(x), int(y), img7, gt, exp["val_preds"][idx], exp["val_probs"][idx])
+        pred  = exp["val_preds"][idx] if (exp and idx < len(exp["val_preds"])) else None
+        probs = exp["val_probs"][idx] if (exp and idx < len(exp["val_probs"])) else None
+        return pixel_info_markdown(int(x), int(y), img7, gt, pred, probs)
     except Exception as e:
-        return f"### Click error: `{type(e).__name__}: {e}`"
-
-
-def handle_click_exp_a(evt, dataset_state, experiments, sel_a, sample_idx):
-    return _handle_click_experiment(evt, dataset_state, experiments, sel_a, sample_idx)
-
-def handle_click_exp_b(evt, dataset_state, experiments, sel_b, sample_idx):
-    return _handle_click_experiment(evt, dataset_state, experiments, sel_b, sample_idx)
+        return f"Click error: `{e}`"

visualize.py

@@ -1,6 +1,15 @@
+import io
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
 import numpy as np
-from config import CLASS_COLORS, IGNORE_INDEX
+from PIL import Image, ImageDraw
+
+from config import CLASS_COLORS, CLASS_NAMES, BAND_NAMES, BAND_DESCRIPTIONS, IGNORE_INDEX, NUM_CLASSES
+
 
+# ── Low-level helpers ─────────────────────────────────────────
 
 def percentile_stretch(x: np.ndarray, low: float = 2.0, high: float = 98.0) -> np.ndarray:
     x = x.astype(np.float32)
@@ -11,16 +20,80 @@ def percentile_stretch(x: np.ndarray, low: float = 2.0, high: float = 98.0) -> n
     return np.clip((x - lo) / (hi - lo), 0, 1)
 
 
+def _fig_to_numpy(fig) -> np.ndarray:
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight", dpi=110)
+    plt.close(fig)
+    buf.seek(0)
+    return np.array(Image.open(buf).convert("RGB"))
+
+
+def _blank_rgb(h: int = 300, w: int = 400) -> np.ndarray:
+    return np.full((h, w, 3), 220, dtype=np.uint8)
+
+
+# ── Composite rendering (full image or patch) ─────────────────
+
+def render_composite(img7: np.ndarray, r: int, g: int, b: int) -> np.ndarray:
+    """img7: (7, H, W) -> (H, W, 3) uint8."""
+    return (np.stack([
+        percentile_stretch(img7[r]),
+        percentile_stretch(img7[g]),
+        percentile_stretch(img7[b]),
+    ], axis=-1) * 255).astype(np.uint8)
+
+
+def render_single_band(img7: np.ndarray, band_idx: int) -> np.ndarray:
+    """Single band as grayscale RGB."""
+    gray = (percentile_stretch(img7[band_idx]) * 255).astype(np.uint8)
+    return np.stack([gray, gray, gray], axis=-1)
+
+
 def multispectral_to_rgb(img7: np.ndarray) -> np.ndarray:
-    """img7: (7, H, W) — uses H_3/H_2/H_1 for natural colour-like composite."""
-    r = percentile_stretch(img7[2])
-    g = percentile_stretch(img7[1])
-    b = percentile_stretch(img7[0])
-    return (np.stack([r, g, b], axis=-1) * 255).astype(np.uint8)
+    """Natural colour composite: H_4/H_3/H_2 -> R/G/B."""
+    return render_composite(img7, r=3, g=2, b=1)
+
+
+# ── Label markers on full-scene image ────────────────────────
+
+def add_labels_overlay(
+    base_rgb:   np.ndarray,
+    train_mask: np.ndarray,
+    val_mask:   np.ndarray,
+    radius:     int = 5,
+) -> np.ndarray:
+    """
+    Draw class-coloured markers on base_rgb.
+    Training labels -> filled squares; validation labels -> circles with white ring.
+    """
+    img  = Image.fromarray(base_rgb)
+    draw = ImageDraw.Draw(img)
+    H, W = base_rgb.shape[:2]
+
+    for cls_idx in range(NUM_CLASSES):
+        color = tuple(int(c) for c in CLASS_COLORS[cls_idx])
+
+        ys, xs = np.where(train_mask == cls_idx)
+        for y, x in zip(ys.tolist(), xs.tolist()):
+            box = [x - radius, y - radius, x + radius, y + radius]
+            box = [max(0, box[0]), max(0, box[1]), min(W-1, box[2]), min(H-1, box[3])]
+            draw.rectangle(box, fill=color, outline=(255, 255, 255))
+
+        ys, xs = np.where(val_mask == cls_idx)
+        for y, x in zip(ys.tolist(), xs.tolist()):
+            outer = [x-radius-2, y-radius-2, x+radius+2, y+radius+2]
+            inner = [x-radius,   y-radius,   x+radius,   y+radius  ]
+            outer = [max(0, v) for v in outer]
+            draw.ellipse(outer, fill=(255, 255, 255))
+            draw.ellipse(inner, fill=color)
 
+    return np.array(img)
+
+
+# ── Mask colourisation ────────────────────────────────────────
 
 def mask_to_color(mask: np.ndarray) -> np.ndarray:
-    """Class indices → RGB. IGNORE_INDEX pixels rendered as light gray."""
+    """Class indices -> RGB. IGNORE_INDEX pixels rendered as light gray."""
     out = np.full((*mask.shape, 3), 200, dtype=np.uint8)
     labeled = (mask != IGNORE_INDEX) & (mask >= 0)
     if labeled.any():
@@ -35,7 +108,7 @@ def overlay_mask(rgb: np.ndarray, mask: np.ndarray, alpha: float = 0.45) -> np.n
 
 
 def correctness_map(pred: np.ndarray, gt: np.ndarray) -> np.ndarray:
-    """Green = correct, red = wrong, gray = unlabeled (IGNORE_INDEX)."""
+    """Green = correct, red = wrong, gray = unlabeled."""
     out = np.full((*pred.shape, 3), 180, dtype=np.uint8)
     labeled = gt != IGNORE_INDEX
     out[labeled & (pred == gt)] = [0, 220, 0]
@@ -47,3 +120,90 @@ def correctness_overlay(rgb: np.ndarray, pred: np.ndarray, gt: np.ndarray, alpha
     cm = correctness_map(pred, gt)
     out = ((1 - alpha) * rgb.astype(np.float32) + alpha * cm.astype(np.float32)).clip(0, 255)
     return out.astype(np.uint8)
+
+
+# ── Full-scene prediction rendering ──────────────────────────
+
+def render_full_prediction_overlay(
+    full_image: np.ndarray,
+    full_pred:  np.ndarray,
+    val_mask:   np.ndarray,
+    alpha:      float = 0.40,
+    dot_radius: int   = 6,
+) -> np.ndarray:
+    """
+    Blend predicted class colours over natural-colour composite, then draw
+    correctness markers at every validation label location.
+    """
+    rgb  = render_composite(full_image, r=3, g=2, b=1)
+    base = overlay_mask(rgb, full_pred, alpha=alpha)
+
+    img  = Image.fromarray(base)
+    draw = ImageDraw.Draw(img)
+    H, W = base.shape[:2]
+
+    for cls_idx in range(NUM_CLASSES):
+        ys, xs = np.where(val_mask == cls_idx)
+        for y, x in zip(ys.tolist(), xs.tolist()):
+            correct = (full_pred[y, x] == cls_idx)
+            ring    = (0, 200, 0) if correct else (220, 0, 0)
+            r = dot_radius
+            outer = [max(0, x-r-2), max(0, y-r-2), min(W-1, x+r+2), min(H-1, y+r+2)]
+            inner = [max(0, x-r),   max(0, y-r),   min(W-1, x+r),   min(H-1, y+r)  ]
+            draw.ellipse(outer, fill=(255, 255, 255))
+            draw.ellipse(inner, fill=ring)
+
+    return np.array(img)
+
+
+# ── Matplotlib charts ─────────────────────────────────────────
+
+def render_spectral_signatures_chart(signatures: dict) -> np.ndarray:
+    """Line chart of per-class mean ± 1-sigma across the 7 bands."""
+    fig, ax = plt.subplots(figsize=(8, 3.8))
+    x = np.arange(len(BAND_NAMES))
+
+    for cls_idx, sig in signatures.items():
+        mean  = sig["mean"]
+        std   = sig["std"]
+        n     = sig["n"]
+        color = CLASS_COLORS[cls_idx] / 255.0
+        label = f"{CLASS_NAMES[cls_idx]} (n={n})"
+        ax.plot(x, mean, "o-", color=color, label=label, linewidth=2, markersize=5)
+        ax.fill_between(x, mean - std, mean + std, alpha=0.18, color=color)
+
+    ax.set_xticks(x)
+    ax.set_xticklabels([d.replace(" (", "\n(") for d in BAND_DESCRIPTIONS], fontsize=8)
+    ax.set_ylabel("Normalised reflectance")
+    ax.set_title("Spectral Signatures by Land Cover Class")
+    ax.legend(loc="upper left", fontsize=8)
+    ax.grid(True, alpha=0.3)
+    fig.tight_layout()
+    return _fig_to_numpy(fig)
+
+
+def render_index_map(
+    index_arr:  np.ndarray,
+    name:       str,
+    train_mask: np.ndarray,
+    val_mask:   np.ndarray,
+) -> np.ndarray:
+    """NDVI or NDWI heatmap with class-coloured label markers."""
+    cmap = "RdYlGn" if name == "NDVI" else "RdYlBu"
+    fig, ax = plt.subplots(figsize=(10, 4.5))
+    im = ax.imshow(index_arr, cmap=cmap, vmin=-1, vmax=1, aspect="auto")
+    plt.colorbar(im, ax=ax, fraction=0.018, pad=0.02)
+
+    for cls_idx in range(NUM_CLASSES):
+        color = CLASS_COLORS[cls_idx] / 255.0
+        name_ = CLASS_NAMES[cls_idx]
+        ys, xs = np.where(train_mask == cls_idx)
+        ax.scatter(xs, ys, c=[color], s=18, marker="s", label=f"{name_} (train)", zorder=5)
+        ys, xs = np.where(val_mask == cls_idx)
+        ax.scatter(xs, ys, c=[color], s=18, marker="o",
+                   edgecolors="white", linewidths=0.6, zorder=6)
+
+    ax.set_title(f"{name} — squares=training labels, circles=val labels")
+    ax.legend(loc="upper right", fontsize=7, markerscale=1.4)
+    fig.tight_layout()
+    return _fig_to_numpy(fig)