Upload 7 files

app.py

@@ -0,0 +1,176 @@
+import gradio as gr
+
+from config import APP_TITLE, set_seed, SEED
+from train import (
+    load_dataset_action,
+    update_explorer_sample,
+    update_compare_sample,
+    train_experiment,
+    handle_click_dataset,
+    handle_click_exp_a,
+    handle_click_exp_b,
+    handle_click_exp_c,
+)
+
+set_seed(SEED)
+
+custom_css = """
+#compare-a img, #compare-b img, #compare-c img, #explorer img {
+    image-rendering: pixelated;
+}
+.small-note { font-size: 0.9rem; opacity: 0.85; }
+"""
+
+with gr.Blocks(title=APP_TITLE, css=custom_css) as demo:
+    gr.Markdown(f"# {APP_TITLE}\nInteractive teaching app for multispectral semantic segmentation.")
+
+    dataset_state = gr.State(None)
+    experiments_state = gr.State([])
+
+    # ── Tab 1: Image Explorer ────────────────────────────────
+    with gr.Tab("1) Image explorer"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                train_size = gr.Slider(60, 2000, value=240, step=20, label="Train subset size")
+                val_size   = gr.Slider(20, 500,  value=60,  step=10, label="Validation subset size")
+                image_size = gr.Slider(64, 256,  value=128, step=32, label="Image size")
+                load_btn   = gr.Button("Load / rebuild dataset", variant="primary")
+                dataset_info = gr.Markdown("### No dataset loaded yet")
+                gr.Markdown(
+                    "<div class='small-note'>Uses procedural synthetic data. "
+                    "See <code>data.py → load_data()</code> to plug in a real dataset.</div>"
+                )
+
+            with gr.Column(scale=2, elem_id="explorer"):
+                explorer_sample_index = gr.Slider(0, 59, value=0, step=1, label="Validation sample index")
+                with gr.Row():
+                    explorer_rgb     = gr.Image(label="RGB / false-color",   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.Row():
+            with gr.Column(scale=1):
+                run_name      = gr.Textbox(label="Experiment name", placeholder="e.g. lr-1e-3_ep-5")
+                slot_label    = gr.Radio(choices=["A", "B", "C"], value="A", label="Save to slot")
+                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")
+            with gr.Column(scale=1):
+                train_summary = gr.Markdown("### No training run yet")
+                gr.Markdown(
+                    "<div class='small-note'>Each slot (A / B / C) stores one run independently. "
+                    "Overwrite a slot to update it. Results appear in the <b>Result comparison</b> tab.</div>"
+                )
+
+    # ── Tab 3: Result Comparison ─────────────────────────────
+    with gr.Tab("3) Result comparison"):
+        compare_sample_index = gr.Slider(0, 59, value=0, step=1, label="Validation sample index")
+        with gr.Row():
+            with gr.Column(scale=1, elem_id="compare-a"):
+                gr.Markdown("## Slot A")
+                compare_a_rgb     = gr.Image(label="Reference RGB",      type="numpy", height=380)
+                compare_a_pred    = gr.Image(label="Prediction mask",     type="numpy", height=380)
+                compare_a_overlay = gr.Image(label="Prediction overlay",  type="numpy", height=380)
+                compare_a_metrics = gr.Markdown("### No experiment")
+                compare_a_error   = gr.Image(label="Correctness map",     type="numpy", height=380)
+                compare_a_click   = gr.Markdown("### Click overlay to inspect pixel")
+
+            with gr.Column(scale=1, elem_id="compare-b"):
+                gr.Markdown("## Slot B")
+                compare_b_rgb     = gr.Image(label="Reference RGB",      type="numpy", height=380)
+                compare_b_pred    = gr.Image(label="Prediction mask",     type="numpy", height=380)
+                compare_b_overlay = gr.Image(label="Prediction overlay",  type="numpy", height=380)
+                compare_b_metrics = gr.Markdown("### No experiment")
+                compare_b_error   = gr.Image(label="Correctness map",     type="numpy", height=380)
+                compare_b_click   = gr.Markdown("### Click overlay to inspect pixel")
+
+            with gr.Column(scale=1, elem_id="compare-c"):
+                gr.Markdown("## Slot C")
+                compare_c_rgb     = gr.Image(label="Reference RGB",      type="numpy", height=380)
+                compare_c_pred    = gr.Image(label="Prediction mask",     type="numpy", height=380)
+                compare_c_overlay = gr.Image(label="Prediction overlay",  type="numpy", height=380)
+                compare_c_metrics = gr.Markdown("### No experiment")
+                compare_c_error   = gr.Image(label="Correctness map",     type="numpy", height=380)
+                compare_c_click   = gr.Markdown("### Click overlay to inspect pixel")
+
+    # ── Shared output lists ───────────────────────────────────
+    _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,
+        compare_c_rgb, compare_c_pred, compare_c_overlay, compare_c_metrics, compare_c_error, compare_c_click,
+    ]
+
+    # ── Event bindings ────────────────────────────────────────
+
+    # Load dataset → reset experiments, update explorer, reset compare slider
+    load_btn.click(
+        fn=load_dataset_action,
+        inputs=[train_size, val_size, image_size],
+        outputs=[
+            dataset_state,
+            experiments_state,
+            dataset_info,
+            explorer_rgb, explorer_gt, explorer_overlay,
+            explorer_click_info,
+            explorer_sample_index,
+            compare_sample_index,
+        ],
+    )
+
+    # Explorer sample slider → update Tab 1 images
+    explorer_sample_index.change(
+        fn=update_explorer_sample,
+        inputs=[dataset_state, explorer_sample_index],
+        outputs=[explorer_rgb, explorer_gt, explorer_overlay, explorer_click_info],
+    )
+
+    # Click on explorer image → pixel info
+    explorer_rgb.select(
+        fn=handle_click_dataset,
+        inputs=[dataset_state, explorer_sample_index],
+        outputs=[explorer_click_info],
+    )
+
+    # Train → update experiments + Tab 3
+    train_btn.click(
+        fn=train_experiment,
+        inputs=[
+            dataset_state, experiments_state,
+            slot_label, learning_rate, batch_size, epochs, base_channels,
+            run_name,
+        ],
+        outputs=[experiments_state, train_summary, compare_sample_index, *_compare_outputs],
+    )
+
+    # Compare sample slider → update Tab 3
+    compare_sample_index.change(
+        fn=update_compare_sample,
+        inputs=[dataset_state, experiments_state, compare_sample_index],
+        outputs=_compare_outputs,
+    )
+
+    # Click on overlay images → pixel info
+    compare_a_overlay.select(
+        fn=handle_click_exp_a,
+        inputs=[dataset_state, experiments_state, compare_sample_index],
+        outputs=[compare_a_click],
+    )
+    compare_b_overlay.select(
+        fn=handle_click_exp_b,
+        inputs=[dataset_state, experiments_state, compare_sample_index],
+        outputs=[compare_b_click],
+    )
+    compare_c_overlay.select(
+        fn=handle_click_exp_c,
+        inputs=[dataset_state, experiments_state, compare_sample_index],
+        outputs=[compare_c_click],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

config.py

@@ -0,0 +1,40 @@
+import numpy as np
+import torch
+
+APP_TITLE = "Multispectral Segmentation Lab"
+SEED = 42
+DEFAULT_IMAGE_SIZE = 128
+NUM_CHANNELS = 7
+NUM_CLASSES = 8
+BAND_NAMES = ["B02", "B03", "B04", "B05", "B06", "B08", "B11"]
+CLASS_NAMES = [
+    "Forest",
+    "Shrubland",
+    "Grassland",
+    "Wetland",
+    "Cropland",
+    "Urban/Built-up",
+    "Barren",
+    "Water",
+]
+CLASS_COLORS = np.array(
+    [
+        [34, 139, 34],    # Forest
+        [154, 205, 50],   # Shrubland
+        [124, 252, 0],    # Grassland
+        [0, 128, 128],    # Wetland
+        [255, 215, 0],    # Cropland
+        [178, 34, 34],    # Urban/Built-up
+        [210, 180, 140],  # Barren
+        [30, 144, 255],   # Water
+    ],
+    dtype=np.uint8,
+)
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def set_seed(seed: int = SEED):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)

data.py

@@ -0,0 +1,140 @@
+import math
+from typing import Tuple, Dict, Optional
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from config import SEED, DEFAULT_IMAGE_SIZE, NUM_CHANNELS, NUM_CLASSES
+
+
+def _draw_disk(mask: np.ndarray, center_y: int, center_x: int, radius: int, value: int):
+    h, w = mask.shape
+    yy, xx = np.ogrid[:h, :w]
+    mask[(yy - center_y) ** 2 + (xx - center_x) ** 2 <= radius ** 2] = value
+
+
+def _draw_rect(mask: np.ndarray, y0: int, x0: int, y1: int, x1: int, value: int):
+    mask[max(0, y0):min(mask.shape[0], y1), max(0, x0):min(mask.shape[1], x1)] = value
+
+
+def generate_synthetic_sample(size: int = DEFAULT_IMAGE_SIZE, seed: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
+    rng = np.random.default_rng(seed)
+    h = w = size
+
+    mask = np.full((h, w), 2, dtype=np.int64)
+
+    for _ in range(rng.integers(1, 3)):
+        cy, cx = rng.integers(h // 6, 5 * h // 6, size=2)
+        _draw_disk(mask, int(cy), int(cx), int(rng.integers(h // 10, h // 5)), 7)
+
+    for cls in [0, 1, 0, 1]:
+        cy, cx = rng.integers(h // 8, 7 * h // 8, size=2)
+        _draw_disk(mask, int(cy), int(cx), int(rng.integers(h // 12, h // 6)), cls)
+
+    water = mask == 7
+    wet = np.zeros_like(water)
+    for dy in [-1, 0, 1]:
+        for dx in [-1, 0, 1]:
+            wet |= np.roll(np.roll(water, dy, axis=0), dx, axis=1)
+    wet &= ~water
+    mask[wet & (rng.random((h, w)) > 0.25)] = 3
+
+    for _ in range(rng.integers(1, 3)):
+        y0 = int(rng.integers(0, h - h // 4))
+        x0 = int(rng.integers(0, w - w // 4))
+        hh = int(rng.integers(h // 8, h // 4))
+        ww = int(rng.integers(w // 8, w // 3))
+        _draw_rect(mask, y0, x0, y0 + hh, x0 + ww, 4)
+        for row in range(y0, min(h, y0 + hh), 6):
+            mask[row: min(h, row + 2), x0: min(w, x0 + ww)] = 2
+
+    for _ in range(rng.integers(1, 4)):
+        y0 = int(rng.integers(0, h - h // 5))
+        x0 = int(rng.integers(0, w - w // 5))
+        _draw_rect(mask, y0, x0, y0 + int(rng.integers(h // 10, h // 5)), x0 + int(rng.integers(w // 10, w // 5)), 5)
+
+    if rng.random() > 0.3:
+        road_y = int(rng.integers(h // 5, 4 * h // 5))
+        mask[max(0, road_y - 1):min(h, road_y + 2), :] = 5
+    if rng.random() > 0.5:
+        road_x = int(rng.integers(w // 5, 4 * w // 5))
+        mask[:, max(0, road_x - 1):min(w, road_x + 2)] = 5
+
+    for _ in range(rng.integers(1, 3)):
+        cy, cx = rng.integers(h // 8, 7 * h // 8, size=2)
+        _draw_disk(mask, int(cy), int(cx), int(rng.integers(h // 14, h // 8)), 6)
+
+    signatures = np.array([
+        [0.10, 0.14, 0.10, 0.25, 0.36, 0.60, 0.24],  # Forest
+        [0.13, 0.18, 0.14, 0.24, 0.30, 0.47, 0.23],  # Shrubland
+        [0.16, 0.22, 0.17, 0.26, 0.32, 0.50, 0.20],  # Grassland
+        [0.09, 0.13, 0.11, 0.18, 0.22, 0.30, 0.10],  # Wetland
+        [0.18, 0.24, 0.20, 0.30, 0.36, 0.52, 0.18],  # Cropland
+        [0.24, 0.26, 0.28, 0.30, 0.31, 0.33, 0.36],  # Urban
+        [0.28, 0.30, 0.32, 0.34, 0.35, 0.36, 0.38],  # Barren
+        [0.05, 0.04, 0.03, 0.02, 0.02, 0.01, 0.00],  # Water
+    ], dtype=np.float32)
+
+    img = np.zeros((NUM_CHANNELS, h, w), dtype=np.float32)
+    for c in range(NUM_CLASSES):
+        region = mask == c
+        for b in range(NUM_CHANNELS):
+            img[b][region] = signatures[c, b]
+
+    yy, xx = np.mgrid[0:h, 0:w]
+    grad1 = (xx / max(1, w - 1)).astype(np.float32)
+    grad2 = (yy / max(1, h - 1)).astype(np.float32)
+    for b in range(NUM_CHANNELS):
+        img[b] += 0.03 * np.sin((b + 1) * grad1 * math.pi)
+        img[b] += 0.02 * np.cos((b + 2) * grad2 * math.pi)
+        img[b] += rng.normal(0, 0.02, size=(h, w)).astype(np.float32)
+
+    return np.clip(img, 0.0, 1.0), mask
+
+
+class MultiSpectralDataset(Dataset):
+    def __init__(self, images: np.ndarray, masks: np.ndarray):
+        self.images = images.astype(np.float32)
+        self.masks = masks.astype(np.int64)
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, idx: int):
+        return torch.from_numpy(self.images[idx]), torch.from_numpy(self.masks[idx])
+
+
+def build_synthetic_dataset(
+    train_size: int, val_size: int, image_size: int
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, str]:
+    total = train_size + val_size
+    images, masks = [], []
+    for i in range(total):
+        img, mask = generate_synthetic_sample(size=image_size, seed=SEED + i)
+        images.append(img)
+        masks.append(mask)
+    images = np.stack(images)
+    masks = np.stack(masks)
+    status = f"Synthetic data | Train: {train_size} | Val: {val_size} | Size: {image_size}×{image_size}"
+    return images[:train_size], masks[:train_size], images[train_size:], masks[train_size:], status
+
+
+def load_data(train_size: int, val_size: int, image_size: int) -> Dict[str, object]:
+    """
+    Load dataset. Currently uses procedural synthetic data.
+
+    TODO: To plug in your own real dataset, replace the call below with a
+    custom loader that returns numpy arrays:
+        - images: (N, 7, H, W)  float32, values in [0, 1]
+        - masks:  (N, H, W)     int64,   class indices in [0, NUM_CLASSES)
+    Then assign tr_x, tr_y, va_x, va_y accordingly and update `status`.
+    """
+    tr_x, tr_y, va_x, va_y, status = build_synthetic_dataset(train_size, val_size, image_size)
+    return {
+        "train_images": tr_x,
+        "train_masks": tr_y,
+        "val_images": va_x,
+        "val_masks": va_y,
+        "status": status,
+    }

metrics.py

@@ -0,0 +1,49 @@
+from typing import Dict, Optional
+import numpy as np
+from config import NUM_CLASSES, CLASS_NAMES
+
+
+def compute_metrics(pred: np.ndarray, gt: np.ndarray, num_classes: int = NUM_CLASSES) -> Dict[str, object]:
+    pred = pred.astype(np.int64)
+    gt = gt.astype(np.int64)
+    cm = np.zeros((num_classes, num_classes), dtype=np.int64)
+    flat_gt = gt.reshape(-1)
+    flat_pred = pred.reshape(-1)
+    for g, p in zip(flat_gt, flat_pred):
+        if 0 <= g < num_classes and 0 <= p < num_classes:
+            cm[g, p] += 1
+
+    overall_acc = float((flat_gt == flat_pred).mean())
+    per_class_acc = []
+    per_class_iou = []
+    for c in range(num_classes):
+        tp = cm[c, c]
+        gt_total = cm[c, :].sum()
+        pred_total = cm[:, c].sum()
+        union = gt_total + pred_total - tp
+        acc = float(tp / gt_total) if gt_total > 0 else None
+        iou = float(tp / union) if union > 0 else None
+        per_class_acc.append(acc)
+        per_class_iou.append(iou)
+    miou = float(np.nanmean([x if x is not None else np.nan for x in per_class_iou]))
+    return {
+        "overall_acc": overall_acc,
+        "miou": miou,
+        "per_class_acc": per_class_acc,
+        "per_class_iou": per_class_iou,
+        "confusion_matrix": cm.tolist(),
+    }
+
+
+def metrics_markdown(metrics: Dict[str, object], title: str = "Metrics") -> str:
+    lines = [f"### {title}"]
+    lines.append(f"- Overall accuracy: **{metrics['overall_acc'] * 100:.2f}%**")
+    lines.append(f"- Mean IoU: **{metrics['miou'] * 100:.2f}%**")
+    lines.append("")
+    lines.append("| Class | Accuracy | IoU |")
+    lines.append("|---|---:|---:|")
+    for name, acc, iou in zip(CLASS_NAMES, metrics["per_class_acc"], metrics["per_class_iou"]):
+        acc_s = "—" if acc is None else f"{acc * 100:.1f}%"
+        iou_s = "—" if iou is None else f"{iou * 100:.1f}%"
+        lines.append(f"| {name} | {acc_s} | {iou_s} |")
+    return "\n".join(lines)

model.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+from config import NUM_CHANNELS, NUM_CLASSES
+
+
+class DoubleConv(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_ch),
+            nn.ReLU(inplace=True),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+class SmallUNet(nn.Module):
+    def __init__(self, in_channels: int = NUM_CHANNELS, num_classes: int = NUM_CLASSES, base_channels: int = 16):
+        super().__init__()
+        self.enc1 = DoubleConv(in_channels, base_channels)
+        self.pool1 = nn.MaxPool2d(2)
+        self.enc2 = DoubleConv(base_channels, base_channels * 2)
+        self.pool2 = nn.MaxPool2d(2)
+        self.enc3 = DoubleConv(base_channels * 2, base_channels * 4)
+        self.pool3 = nn.MaxPool2d(2)
+
+        self.bottleneck = DoubleConv(base_channels * 4, base_channels * 8)
+
+        self.up3 = nn.ConvTranspose2d(base_channels * 8, base_channels * 4, kernel_size=2, stride=2)
+        self.dec3 = DoubleConv(base_channels * 8, base_channels * 4)
+        self.up2 = nn.ConvTranspose2d(base_channels * 4, base_channels * 2, kernel_size=2, stride=2)
+        self.dec2 = DoubleConv(base_channels * 4, base_channels * 2)
+        self.up1 = nn.ConvTranspose2d(base_channels * 2, base_channels, kernel_size=2, stride=2)
+        self.dec1 = DoubleConv(base_channels * 2, base_channels)
+
+        self.head = nn.Conv2d(base_channels, num_classes, kernel_size=1)
+
+    def forward(self, x):
+        e1 = self.enc1(x)
+        e2 = self.enc2(self.pool1(e1))
+        e3 = self.enc3(self.pool2(e2))
+        b = self.bottleneck(self.pool3(e3))
+
+        d3 = self.up3(b)
+        d3 = torch.cat([d3, e3], dim=1)
+        d3 = self.dec3(d3)
+
+        d2 = self.up2(d3)
+        d2 = torch.cat([d2, e2], dim=1)
+        d2 = self.dec2(d2)
+
+        d1 = self.up1(d2)
+        d1 = torch.cat([d1, e1], dim=1)
+        d1 = self.dec1(d1)
+
+        return self.head(d1)

train.py

@@ -0,0 +1,294 @@
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import gradio as gr
+from PIL import Image
+
+from config import DEVICE, NUM_CHANNELS, NUM_CLASSES, DEFAULT_IMAGE_SIZE, BAND_NAMES, CLASS_NAMES
+from data import MultiSpectralDataset, load_data
+from model import SmallUNet
+from visualize import multispectral_to_rgb, mask_to_color, overlay_mask, correctness_overlay
+from metrics import compute_metrics, metrics_markdown
+
+
+# ── Inference ────────────────────────────────────────────────
+
+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)
+    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)
+    preds, probs = [], []
+    model.eval()
+    with torch.no_grad():
+        for xb, _ in loader:
+            xb = xb.to(DEVICE)
+            pb = F.softmax(model(xb), dim=1)
+            preds.append(torch.argmax(pb, dim=1).cpu().numpy())
+            probs.append(pb.cpu().numpy())
+    return np.concatenate(preds, axis=0), np.concatenate(probs, axis=0)
+
+
+# ── Render helpers ───────────────────────────────────────────
+
+def _blank(size: int = DEFAULT_IMAGE_SIZE) -> Image.Image:
+    return Image.fromarray(np.full((size, size, 3), 245, dtype=np.uint8))
+
+
+def pixel_info_markdown(
+    x: int, y: int,
+    img7: np.ndarray, gt: np.ndarray,
+    pred: Optional[np.ndarray], probs: Optional[np.ndarray],
+) -> str:
+    h, w = gt.shape
+    x = int(np.clip(x, 0, w - 1))
+    y = int(np.clip(y, 0, h - 1))
+    lines = [f"### Pixel ({x}, {y})", f"- Ground truth: **{CLASS_NAMES[int(gt[y, x])]}**"]
+    if pred is not None:
+        pred_class = int(pred[y, x])
+        lines.append(f"- Prediction: **{CLASS_NAMES[pred_class]}**")
+        lines.append(f"- Correct: **{'Yes' if pred_class == int(gt[y, x]) 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: —")
+    lines += ["", "**Band values**"] + [f"- {n}: {float(img7[b, y, x]):.3f}" for b, n in enumerate(BAND_NAMES)]
+    return "\n".join(lines)
+
+
+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, "### Click info")
+    if dataset_state is None or "val_images" not in dataset_state:
+        return no_data
+    val_images = dataset_state["val_images"]
+    val_masks = dataset_state["val_masks"]
+    if len(val_images) == 0:
+        return no_data
+
+    idx = max(0, min(int(sample_idx), len(val_images) - 1))
+    rgb = multispectral_to_rgb(val_images[idx])
+    gt = val_masks[idx]
+
+    if exp is None:
+        return (
+            rgb, mask_to_color(gt), overlay_mask(rgb, gt),
+            "### No experiment selected",
+            _blank(),
+            pixel_info_markdown(0, 0, val_images[idx], gt, None, None),
+        )
+
+    # Guard: experiment predictions might be from a different dataset
+    if idx >= len(exp["val_preds"]):
+        return (
+            rgb, mask_to_color(gt), overlay_mask(rgb, gt),
+            "### Dataset reloaded — retrain to refresh",
+            _blank(),
+            "### Retrain needed",
+        )
+
+    pred = exp["val_preds"][idx].astype(np.uint8)
+    probs = exp["val_probs"][idx].astype(np.float32)
+    sample_metrics = compute_metrics(pred, gt, num_classes=NUM_CLASSES)
+    return (
+        rgb,
+        mask_to_color(pred),
+        overlay_mask(rgb, pred),
+        metrics_markdown(sample_metrics, title=f"Slot {exp['slot']} — {exp['name']} (sample {idx})"),
+        correctness_overlay(rgb, pred, gt),
+        pixel_info_markdown(0, 0, val_images[idx], gt, pred, probs),
+    )
+
+
+def render_compare_view(dataset_state, experiments, sample_idx: int) -> Tuple:
+    """Returns 18 values: 6 outputs × 3 slots (A, B, C)."""
+    slot_map = {e["slot"]: e for e in experiments}
+    return (
+        *render_experiment_panel(dataset_state, slot_map.get("A"), sample_idx),
+        *render_experiment_panel(dataset_state, slot_map.get("B"), sample_idx),
+        *render_experiment_panel(dataset_state, slot_map.get("C"), sample_idx),
+    )
+
+
+# ── Gradio action functions ──────────────────────────────────
+
+def load_dataset_action(train_size: int, val_size: int, image_size: int):
+    """
+    Loads a fresh dataset and resets all experiment state.
+    Returns 9 values for Gradio outputs.
+    """
+    train_size, val_size, image_size = int(train_size), int(val_size), int(image_size)
+    dataset_state = load_data(train_size, val_size, image_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 (synthetic)",
+        f"- {dataset_state['status']}",
+        f"- Channels: **{NUM_CHANNELS}** ({', '.join(BAND_NAMES)})",
+        f"- Classes: **{NUM_CLASSES}** ({', '.join(CLASS_NAMES)})",
+        "",
+        "_Using procedural synthetic data. See `data.py → load_data()` to plug in a real dataset._",
+    ])
+
+    return (
+        dataset_state,
+        [],                                                                  # reset experiments_state
+        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
+    )
+
+
+def update_explorer_sample(dataset_state, sample_idx: int):
+    """Updates the Tab 1 explorer images when the sample index slider changes."""
+    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_compare_sample(dataset_state, experiments, sample_idx: int):
+    """Updates Tab 3 when the compare sample index slider changes."""
+    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, int(sample_idx))
+
+
+def train_experiment(
+    dataset_state: Dict,
+    experiments: List[Dict],
+    slot_label: str,
+    learning_rate: float,
+    batch_size: int,
+    epochs: int,
+    base_channels: int,
+    run_name: str,
+):
+    """
+    Trains a SmallUNet and stores results in the given slot.
+    Returns 21 values: experiments, summary, compare_sample_index update, + 18 compare outputs.
+    """
+    if dataset_state is None or "train_images" not in dataset_state:
+        raise gr.Error("Load a dataset first.")
+
+    train_images = dataset_state["train_images"]
+    train_masks = dataset_state["train_masks"]
+    val_images = dataset_state["val_images"]
+    val_masks = dataset_state["val_masks"]
+
+    loader = DataLoader(
+        MultiSpectralDataset(train_images, train_masks),
+        batch_size=int(batch_size), shuffle=True,
+    )
+    model = SmallUNet(NUM_CHANNELS, NUM_CLASSES, int(base_channels)).to(DEVICE)
+    optimizer = torch.optim.Adam(model.parameters(), lr=float(learning_rate))
+    criterion = nn.CrossEntropyLoss()
+
+    history = []
+    for _ in range(int(epochs)):
+        model.train()
+        total_loss, 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())
+            n += 1
+        history.append(total_loss / max(1, n))
+
+    val_preds, val_probs = build_prediction_cache(model, val_images, batch_size=max(1, int(batch_size)))
+    global_metrics = compute_metrics(val_preds.reshape(-1), val_masks.reshape(-1), num_classes=NUM_CLASSES)
+
+    experiment = {
+        "name": (run_name or f"Run {len(experiments) + 1}").strip(),
+        "slot": slot_label,
+        "config": {
+            "learning_rate": float(learning_rate),
+            "batch_size": int(batch_size),
+            "epochs": int(epochs),
+            "base_channels": int(base_channels),
+        },
+        "train_loss_history": history,
+        "global_metrics": global_metrics,
+        "val_preds": val_preds.astype(np.uint8),
+        "val_probs": val_probs.astype(np.float32),
+    }
+
+    slot_map = {e["slot"]: e for e in experiments}
+    slot_map[slot_label] = experiment
+    experiments = [slot_map[s] for s in ["A", "B", "C"] if s in slot_map]
+
+    summary = "\n".join([
+        f"### Training finished — Slot {slot_label}",
+        f"- Experiment: **{experiment['name']}**",
+        f"- Device: **{DEVICE}** | Epochs: **{int(epochs)}**",
+        f"- Final loss: **{history[-1]:.4f}**",
+        f"- Val accuracy: **{global_metrics['overall_acc'] * 100:.2f}%**",
+        f"- Val mIoU: **{global_metrics['miou'] * 100:.2f}%**",
+    ])
+
+    compare_slider = gr.update(maximum=max(0, len(val_images) - 1), value=0)
+    compare_outputs = render_compare_view(dataset_state, experiments, 0)
+    return experiments, summary, compare_slider, *compare_outputs
+
+
+# ── 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 _handle_click_experiment(evt: gr.SelectData, dataset_state, experiments, slot: str, 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))
+    exp = next((e for e in experiments if e["slot"] == slot), None)
+    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])
+
+
+def handle_click_exp_a(evt, dataset_state, experiments, sample_idx):
+    return _handle_click_experiment(evt, dataset_state, experiments, "A", sample_idx)
+
+def handle_click_exp_b(evt, dataset_state, experiments, sample_idx):
+    return _handle_click_experiment(evt, dataset_state, experiments, "B", sample_idx)
+
+def handle_click_exp_c(evt, dataset_state, experiments, sample_idx):
+    return _handle_click_experiment(evt, dataset_state, experiments, "C", sample_idx)

visualize.py

@@ -0,0 +1,44 @@
+import numpy as np
+from config import CLASS_COLORS
+
+
+def percentile_stretch(x: np.ndarray, low: float = 2.0, high: float = 98.0) -> np.ndarray:
+    x = x.astype(np.float32)
+    lo = np.percentile(x, low)
+    hi = np.percentile(x, high)
+    if hi <= lo:
+        hi = lo + 1e-6
+    return np.clip((x - lo) / (hi - lo), 0, 1)
+
+
+def multispectral_to_rgb(img7: np.ndarray) -> np.ndarray:
+    """img7 shape: (7, H, W) — uses B04, B03, B02 for natural RGB view."""
+    r = percentile_stretch(img7[2])
+    g = percentile_stretch(img7[1])
+    b = percentile_stretch(img7[0])
+    rgb = np.stack([r, g, b], axis=-1)
+    return (rgb * 255).astype(np.uint8)
+
+
+def mask_to_color(mask: np.ndarray) -> np.ndarray:
+    return CLASS_COLORS[mask]
+
+
+def overlay_mask(rgb: np.ndarray, mask: np.ndarray, alpha: float = 0.45) -> np.ndarray:
+    color_mask = mask_to_color(mask)
+    out = ((1 - alpha) * rgb.astype(np.float32) + alpha * color_mask.astype(np.float32)).clip(0, 255)
+    return out.astype(np.uint8)
+
+
+def correctness_map(pred: np.ndarray, gt: np.ndarray) -> np.ndarray:
+    correct = pred == gt
+    out = np.zeros((pred.shape[0], pred.shape[1], 3), dtype=np.uint8)
+    out[correct] = np.array([0, 220, 0], dtype=np.uint8)
+    out[~correct] = np.array([220, 0, 0], dtype=np.uint8)
+    return out
+
+
+def correctness_overlay(rgb: np.ndarray, pred: np.ndarray, gt: np.ndarray, alpha: float = 0.38) -> np.ndarray:
+    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)