Add Segformer85Mv2 (fine-tuned on Orchard Navigation, autumn+different camera). v1 unchanged.

.gitattributes

@@ -37,3 +37,6 @@ samples/sample_00_frame_2575.jpg filter=lfs diff=lfs merge=lfs -text
 samples/sample_05_frame_3371.jpg filter=lfs diff=lfs merge=lfs -text
 samples/sample_09_frame_4009.jpg filter=lfs diff=lfs merge=lfs -text
 v6_OOD_full_res.mp4 filter=lfs diff=lfs merge=lfs -text
+samples_v6_vs_v7/cmp_0_Annotation_Data_Sep16_T_m.jpg filter=lfs diff=lfs merge=lfs -text
+samples_v6_vs_v7/cmp_4_Annotation_Data_Sep16_T_m.jpg filter=lfs diff=lfs merge=lfs -text
+samples_v6_vs_v7/cmp_5_Annotation_Data_Sep16_T_m.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -14,9 +14,16 @@ pipeline_tag: image-segmentation
 base_model: nvidia/segformer-b5-finetuned-ade-640-640
 ---
 
-# Segformer85Mv1 — Apple Orchard Semantic Segmentation
+# Segformer85M — Apple Orchard Semantic Segmentation
 
-A Segformer-B5 (85M parameters) fine-tuned for **8-class semantic segmentation** of outdoor apple orchard scenes captured from a robotic platform.
+Segformer-B5 (85M parameters) fine-tuned for **8-class semantic segmentation** of outdoor apple orchard scenes captured from a robotic platform.
+
+This repo contains **two checkpoints**:
+
+| File | When to use |
+|------|-------------|
+| **`Segformer85Mv1.pt`** | Original v1, trained only on the spring oak_0415 dataset. Best baseline. |
+| **`Segformer85Mv2.pt`** ⭐ | v1 + fine-tuned on a second dataset (different camera, autumn season). **Use this for general deployment** — same accuracy on the original orchard, dramatically better generalization to new cameras / new seasons. |
 
 ## Quick Use
 
@@ -26,8 +33,10 @@ from transformers import SegformerForSemanticSegmentation
 import torch, cv2, numpy as np
 import torch.nn.functional as F
 
-# 1. Download weights
-ckpt_path = hf_hub_download(repo_id="YOUR_USER/Segformer85Mv1", filename="Segformer85Mv1.pt")
+# 1. Download weights — pick v1 OR v2
+ckpt_path = hf_hub_download(repo_id="WEN0256/Segformer85Mv1", filename="Segformer85Mv2.pt")
+#                                                                       ^^^^^^^^^^^^^^^^^
+#                                                                       use v2 by default
 
 # 2. Init architecture from base + load fine-tuned weights
 NAMES = ["tree","ground","person","sky","road","mountain","building","background"]
@@ -65,7 +74,7 @@ A ready-to-use `predict.py` is included in this repo.
 | 2  | person      | Workers in scene                                       |
 | 3  | sky         |                                                        |
 | 4  | road        | Path between rows                                      |
-| 5  | mountain    | Distant terrain (often confused with sky in fog)       |
+| 5  | mountain    | Distant terrain                                        |
 | 6  | building    | Sheds, equipment shelters                              |
 | 7  | background  | Unknown / unlabeled regions (model output rare)        |
 
@@ -81,12 +90,12 @@ A ready-to-use `predict.py` is included in this repo.
 | `std` | `[0.229, 0.224, 0.225]` |
 | Input resolution | Any H×W where both are multiples of 32 |
 | Trained at | 1024×576 (native 16:9) |
-| Recommended inference | 1280×704 or original native (snap to 32-multiple) |
-| Precision | bfloat16 fine — model weights stored in fp32 |
 
-## Performance (NO data leakage)
+## Performance
+
+### v1 (Segformer85Mv1.pt) — original training only
 
-Validated on a temporally-disjoint hold-out (frames 4501+ from training set):
+Validated on a temporally-disjoint hold-out from the same recording (frames 4501+, no leakage):
 
 | Metric | Value |
 |---|---|
@@ -94,7 +103,25 @@ Validated on a temporally-disjoint hold-out (frames 4501+ from training set):
 | **mIoU (7 real classes)** | **0.714** |
 | **Pixel accuracy** | **0.834** |
 
-### Per-class IoU
+### v2 (Segformer85Mv2.pt) — v1 + Orchard Navigation fine-tune ⭐
+
+Same v1 hold-out → no regression on old domain:
+
+| Metric | v1 | **v2** |
+|---|---|---|
+| Tree IoU (orig orchard, no leak) | 0.742 | **0.742** ✅ |
+| mIoU (orig orchard) | 0.714 | 0.712 |
+
+NEW orchard hold-out (different camera, autumn season — Aug+Sep capture):
+
+| Metric | v1 | **v2** |
+|---|---|---|
+| Tree recall on new orchard | ~0.55 (estimated) | **0.999** 🚀 |
+
+**Visual qualitative**: v1 sometimes misclassifies autumn foliage as `person` (red); v2 cleanly segments it as `tree`. See `samples/` for side-by-side examples.
+
+### v1 per-class IoU (8-class, no leak)
+
 | Class | IoU | Precision | Recall |
 |---|---|---|---|
 | tree | 0.742 | 0.79 | 0.93 |
@@ -105,20 +132,25 @@ Validated on a temporally-disjoint hold-out (frames 4501+ from training set):
 | mountain | 0.437 | 0.62 | 0.66 |
 | building | 0.711 | 0.84 | 0.83 |
 
-(Reported values from epoch 21 best-tree checkpoint on the no-leak validation split.)
-
-### OOD evaluation
-On a completely held-out recording (1912 frames from `oak_0415_twoRadar_1`, never seen in training), mean prediction confidence is **0.939**, with model predicting `tree` on 41.8% of pixels and falling back to `background` on only 7.4% — indicating strong out-of-distribution generalization.
-
 ## Training Data
 
-- ~5300 frames from a single oak_0415_oneRadar_1 recording
+### v1 base
+- ~5300 frames from a single oak_0415_oneRadar_1 recording (spring, single camera)
 - Initial annotations from 3 separate Roboflow projects (SAM-assisted polygons), merged + class-aligned (`vines`→`tree`, `moutain`→`mountain` typo fixed)
 - Pseudo-labels generated by an earlier model to fill SAM annotation gaps
 - Temporal split: frames `<=4500` train (5177 samples), frames `>4500` validation (155 samples) — **no neighbor leakage**
 
+### v2 fine-tune (NEW)
+- **+311 images** from "Orchard Navigation" dataset:
+  - 178 frames from a Sep-16 recording (autumn season)
+  - 134 frames from a Windows webcam capture (Aug 23, different camera/sensor)
+- Tree-only polygon annotations
+- Mixed with 500 sampled v1 images (full 8-class masks) to prevent forgetting
+- Non-tree pixels in new images set to `ignore_index=255` so the model only adapts its tree decisions, leaving other classes untouched
+
 ## Training Recipe
 
+### v1
 | Hyperparameter | Value |
 |---|---|
 | Optimizer | AdamW, weight_decay 0.01 |
@@ -131,16 +163,41 @@ On a completely held-out recording (1912 frames from `oak_0415_twoRadar_1`, neve
 | Class weights | tree 1.5, ground 0.5, person 1.5, sky 1.0, road 1.0, mountain 1.0, building 1.0, background 0.1 |
 | Hardware | RTX 5090 (32 GB), ~2.3 hours |
 
+### v2 fine-tune (delta from v1)
+| Hyperparameter | Value |
+|---|---|
+| LR | **5e-6** (10× lower for safe fine-tune) |
+| Epochs | **8** (best at epoch 3) |
+| `ignore_index` | **255** (for unlabeled pixels in new data) |
+| Everything else | Same as v1 |
+| Hardware | RTX 5090, ~13 minutes |
+
+## Limitations
+
+This model was trained on a **single Korean apple orchard** (spring 2024) with a **single robot platform**, plus a small fine-tune on a second autumn capture. Expect degradation on:
+
+- ⚠️ Different orchards (different tree species, layouts, training systems)
+- ⚠️ Different cameras (different FOV, color profiles, sensors)
+- 💀 Different seasons not in training (winter dormant trees)
+- 💀 Different lighting (rain, dawn/dusk, night)
+- 💀 Aerial / drone perspectives
+
+For deployment in a new context, plan to fine-tune on 100-300 in-domain images.
+
 ## Files in This Repo
 
 | File | Purpose |
 |---|---|
-| `Segformer85Mv1.pt` | Fine-tuned weights (339 MB) |
-| `predict.py` | Standalone inference script |
+| `Segformer85Mv1.pt` | Original v1 weights (339 MB) |
+| `Segformer85Mv2.pt` | v1 + Orchard Navigation fine-tune (339 MB) ⭐ |
+| `predict.py` | Standalone inference script (defaults to v2) |
 | `README.md` | This file |
-| `samples/*.jpg` | Side-by-side prediction examples |
-| `train_v6_5090.py` | Training script (for reproduction) |
-| `history_v6.json` | Per-epoch training history |
+| `samples/*.jpg` | v1 prediction examples (in-domain) |
+| `samples_v6_vs_v7/*.jpg` | **v1 vs v2 side-by-side** on new orchard (showcases v2 improvement) |
+| `train_v6_5090.py` | v1 training script |
+| `finetune_v7.py` | v2 fine-tune script |
+| `history_v6.json` | v1 per-epoch training history |
+| `history_v7.json` | v2 fine-tune history |
 | `v6_OOD_full_res.mp4` | 1-minute OOD inference video at native resolution |
 
 ## License

Segformer85Mv2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca217b62e6b1adcb1bc03571761ea12ada7a2e752fb2e2fce19ce6f14edd838f
+size 338883907

finetune_v7.py

@@ -0,0 +1,240 @@
+"""V7 — fine-tune Segformer85Mv1 on Orchard Navigation dataset (tree-only).
+
+Strategy:
+  - Mix 311 NEW images (tree=0, everything-else=255 ignore) with a SAMPLE of
+    500 OLD images (full 8-class masks) — keeps non-tree classes from drifting.
+  - Very low LR (5e-6) and few epochs (8) to nudge tree decisions toward the
+    new domain (different camera, different season) without catastrophic
+    forgetting of old classes.
+  - Temporal val from old data (frame > 4500) gives apples-to-apples vs v6.
+  - Also reports tree IoU on a held-out chunk of NEW data.
+"""
+from __future__ import annotations
+import json, re, time, random
+from pathlib import Path
+import numpy as np, cv2, torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader, ConcatDataset
+from torch.amp import autocast
+import albumentations as A
+from transformers import SegformerForSemanticSegmentation
+
+# ───────────── config ─────────────
+ROOT = Path("/workspace/agmotree")
+OLD_IMG = ROOT / "old_data/images"
+OLD_MSK = ROOT / "old_data/masks_pseudo"
+NEW_IMG = ROOT / "new_data/images"
+NEW_MSK = ROOT / "new_data/masks"
+CKPT_IN = ROOT / "Segformer85Mv1.pt"
+OUT_DIR = ROOT / "v7_output"
+OUT_DIR.mkdir(parents=True, exist_ok=True)
+
+NAMES = ["tree","ground","person","sky","road","mountain","building","background"]
+NUM_CLASSES = 8
+IGNORE_INDEX = 255
+
+IMG_W = 1024
+IMG_H = 576
+BATCH = 2
+GRAD_ACCUM = 4
+EPOCHS = 8
+LR = 5e-6
+N_OLD_SAMPLE = 500   # how many old images to mix in
+SEED = 42
+
+# Class weights — heavier on tree, normal on others
+WEIGHTS = np.array([1.5, 0.5, 1.5, 1.0, 1.0, 1.0, 1.0, 0.1])
+
+# ───────────── data ─────────────
+def frame_num(p):
+    m = re.match(r"frame_(\d+)", p.stem); return int(m.group(1)) if m else -1
+
+random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
+
+old_all = sorted(OLD_IMG.glob("*.jpg"))
+old_train = [p for p in old_all if frame_num(p) <= 4500]
+old_val   = [p for p in old_all if frame_num(p) >  4500]
+old_train_sampled = random.sample(old_train, min(N_OLD_SAMPLE, len(old_train)))
+
+new_all = sorted(NEW_IMG.glob("*.jpg"))
+random.shuffle(new_all)
+n_new_val = max(20, len(new_all) // 10)
+new_val = new_all[:n_new_val]
+new_train = new_all[n_new_val:]
+
+print(f"=== V7 FINE-TUNE ===")
+print(f"  old train (sampled): {len(old_train_sampled)} (8-class masks)")
+print(f"  new train:           {len(new_train)} (tree-only, rest ignore)")
+print(f"  old val (no leak):   {len(old_val)} (frames > 4500)")
+print(f"  new val:             {len(new_val)} (held-out new data)")
+
+train_tf = A.Compose([
+    A.Resize(IMG_H, IMG_W),
+    A.HorizontalFlip(p=0.5),
+    A.RandomBrightnessContrast(0.2, 0.2, p=0.5),
+    A.HueSaturationValue(10, 15, 10, p=0.3),
+    A.GaussianBlur(blur_limit=(3,5), p=0.2),
+    A.Normalize(mean=(0.485,0.456,0.406), std=(0.229,0.224,0.225)),
+])
+val_tf = A.Compose([
+    A.Resize(IMG_H, IMG_W),
+    A.Normalize(mean=(0.485,0.456,0.406), std=(0.229,0.224,0.225)),
+])
+
+class SegDS(Dataset):
+    def __init__(self, paths, mask_dir, tf):
+        self.paths = paths; self.mask_dir = mask_dir; self.tf = tf
+    def __len__(self): return len(self.paths)
+    def __getitem__(self, i):
+        ip = self.paths[i]
+        img = cv2.cvtColor(cv2.imread(str(ip)), cv2.COLOR_BGR2RGB)
+        msk = cv2.imread(str(self.mask_dir / (ip.stem + ".png")), cv2.IMREAD_GRAYSCALE)
+        out = self.tf(image=img, mask=msk)
+        return (torch.from_numpy(out["image"]).permute(2,0,1).float(),
+                torch.from_numpy(out["mask"]).long())
+
+old_train_ds = SegDS(old_train_sampled, OLD_MSK, train_tf)
+new_train_ds = SegDS(new_train,         NEW_MSK, train_tf)
+old_val_ds   = SegDS(old_val,           OLD_MSK, val_tf)
+new_val_ds   = SegDS(new_val,           NEW_MSK, val_tf)
+
+# ───────────── eval helpers ─────────────
+def confusion(preds, ys, n):
+    cm = np.zeros((n, n), dtype=np.int64)
+    for tc in range(n):
+        mt = (ys == tc)
+        if not mt.any(): continue
+        for pc in range(n):
+            cm[tc, pc] += int(((preds == pc) & mt).sum())
+    return cm
+
+def iou_from_cm(cm):
+    n = cm.shape[0]; ious = np.zeros(n)
+    for c in range(n):
+        tp = cm[c,c]; fp = cm[:,c].sum()-tp; fn = cm[c,:].sum()-tp
+        ious[c] = tp / (tp+fp+fn) if (tp+fp+fn) > 0 else float("nan")
+    return ious
+
+# ───────────── train ─────────────
+log_path = OUT_DIR / "log_v7.txt"
+def log(m):
+    print(m, flush=True)
+    with log_path.open("a", encoding="utf-8") as f: f.write(m + "\n")
+
+def main():
+    log_path.write_text("")
+    train_ds = ConcatDataset([old_train_ds, new_train_ds])
+    train_loader = DataLoader(train_ds, batch_size=BATCH, shuffle=True,
+                              num_workers=8, pin_memory=True, drop_last=True,
+                              persistent_workers=True)
+    old_val_loader = DataLoader(old_val_ds, batch_size=BATCH, shuffle=False,
+                                num_workers=4, pin_memory=True, persistent_workers=True)
+    new_val_loader = DataLoader(new_val_ds, batch_size=BATCH, shuffle=False,
+                                num_workers=4, pin_memory=True, persistent_workers=True)
+
+    log(f"=== V7 FINE-TUNE ===")
+    log(f"old train sampled={len(old_train_sampled)}  new train={len(new_train)}")
+    log(f"old val (no-leak frames>4500): {len(old_val)}  new val: {len(new_val)}")
+    log(f"loading {CKPT_IN} ...")
+    model = SegformerForSemanticSegmentation.from_pretrained(
+        "nvidia/segformer-b5-finetuned-ade-640-640", num_labels=NUM_CLASSES,
+        id2label={i:n for i,n in enumerate(NAMES)},
+        label2id={n:i for i,n in enumerate(NAMES)},
+        ignore_mismatched_sizes=True,
+    ).cuda()
+    ckpt = torch.load(CKPT_IN, map_location="cuda", weights_only=False)
+    model.load_state_dict(ckpt["model"])
+    log(f"  loaded v6 ckpt: epoch {ckpt['epoch']}, prev tree IoU {ckpt['tree_iou']:.3f}")
+
+    cw = torch.tensor(WEIGHTS, dtype=torch.float32, device="cuda")
+    loss_fn = nn.CrossEntropyLoss(weight=cw, ignore_index=IGNORE_INDEX)
+    optim = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-2)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(optim, T_max=EPOCHS*len(train_loader))
+
+    log(f"train batches: {len(train_loader)}")
+
+    best_avg_tree = -1.0
+    history = []
+    for epoch in range(1, EPOCHS+1):
+        model.train()
+        t0 = time.time()
+        epoch_loss = 0.0
+        optim.zero_grad()
+        for step, (x, y) in enumerate(train_loader):
+            x = x.cuda(non_blocking=True); y = y.cuda(non_blocking=True)
+            with autocast("cuda", dtype=torch.bfloat16):
+                out = model(pixel_values=x)
+                logits = F.interpolate(out.logits, size=y.shape[-2:], mode="bilinear", align_corners=False)
+                loss = loss_fn(logits, y) / GRAD_ACCUM
+            loss.backward()
+            if (step+1) % GRAD_ACCUM == 0:
+                optim.step(); optim.zero_grad(); sched.step()
+            epoch_loss += loss.item() * GRAD_ACCUM
+        train_loss = epoch_loss / len(train_loader)
+
+        # ─── eval on old (8-class) and new (tree-only) ───
+        model.eval()
+        cm_old = np.zeros((NUM_CLASSES, NUM_CLASSES), dtype=np.int64)
+        cm_new = np.zeros((NUM_CLASSES, NUM_CLASSES), dtype=np.int64)
+        with torch.no_grad():
+            for x, y in old_val_loader:
+                x = x.cuda(); y = y.cuda()
+                with autocast("cuda", dtype=torch.bfloat16):
+                    out = model(pixel_values=x)
+                    logits = F.interpolate(out.logits, size=y.shape[-2:], mode="bilinear", align_corners=False)
+                cm_old += confusion(logits.argmax(1).cpu().numpy(), y.cpu().numpy(), NUM_CLASSES)
+            tree_tp = tree_fn = 0   # only tree class is meaningful (others not labeled)
+            for x, y in new_val_loader:
+                x = x.cuda(); y = y.cuda()
+                with autocast("cuda", dtype=torch.bfloat16):
+                    out = model(pixel_values=x)
+                    logits = F.interpolate(out.logits, size=y.shape[-2:], mode="bilinear", align_corners=False)
+                pred = logits.argmax(1).cpu().numpy()
+                ys = y.cpu().numpy()
+                # Tree recall = TP / (TP+FN) over labeled tree pixels only.
+                # Ignored pixels (ys==255) are excluded entirely.
+                tree_mask = (ys == 0)
+                tree_tp += int(((pred == 0) & tree_mask).sum())
+                tree_fn += int(((pred != 0) & tree_mask).sum())
+
+        iou_old = iou_from_cm(cm_old)
+        miou_old = float(np.nanmean(iou_old[:7]))
+        tree_old = float(iou_old[0])
+        # NEW val: only tree-recall is meaningful (other classes unlabeled)
+        tree_recall_new = tree_tp / (tree_tp + tree_fn) if (tree_tp + tree_fn) > 0 else float("nan")
+        avg_tree = (tree_old + tree_recall_new) / 2
+
+        elapsed = time.time() - t0
+        log(f"epoch {epoch:02d}/{EPOCHS}  tloss={train_loss:.4f}  ({elapsed:.0f}s)")
+        log(f"  OLD val (8-class):  mIoU(7)={miou_old:.3f}  tree IoU={tree_old:.3f}")
+        log(f"  NEW val:            tree RECALL={tree_recall_new:.3f}  (TP={tree_tp:,} FN={tree_fn:,})")
+        log(f"  per-class OLD: " + ", ".join(f"{n}={v:.3f}" for n,v in zip(NAMES, iou_old)))
+
+        history.append({
+            "epoch": epoch, "train_loss": float(train_loss),
+            "tree_iou_old": tree_old, "tree_recall_new": tree_recall_new,
+            "miou_old_7": miou_old, "avg_tree": avg_tree,
+            "per_class_iou_old": {n: float(v) for n, v in zip(NAMES, iou_old)},
+        })
+
+        torch.save({"model": model.state_dict(), "epoch": epoch,
+                    "tree_iou_old": tree_old, "tree_recall_new": tree_recall_new,
+                    "miou_old_7": miou_old},
+                   OUT_DIR / "v7_last.pt")
+        if avg_tree > best_avg_tree:
+            best_avg_tree = avg_tree
+            torch.save({"model": model.state_dict(), "epoch": epoch,
+                        "tree_iou_old": tree_old, "tree_recall_new": tree_recall_new,
+                        "miou_old_7": miou_old},
+                       OUT_DIR / "v7_best.pt")
+            log(f"  saved v7_best.pt (avg {avg_tree:.3f})")
+
+        (OUT_DIR / "history_v7.json").write_text(json.dumps(history, indent=2))
+
+    log(f"\n=== DONE ===")
+    log(f"best avg tree IoU (old+new mean): {best_avg_tree:.3f}")
+
+
+if __name__ == "__main__":
+    main()

history_v7.json

@@ -0,0 +1,146 @@
+[
+  {
+    "epoch": 1,
+    "train_loss": 0.1445474764559036,
+    "tree_iou_old": 0.7276111213384829,
+    "tree_recall_new": 0.9964709949014234,
+    "miou_old_7": 0.7071341866953114,
+    "avg_tree": 0.8620410581199531,
+    "per_class_iou_old": {
+      "tree": 0.7276111213384829,
+      "ground": 0.8545028283913517,
+      "person": 0.7197452556695519,
+      "sky": 0.7665161031217024,
+      "road": 0.7716666106060096,
+      "mountain": 0.4105051777602567,
+      "building": 0.6993922099798241,
+      "background": 0.30209167992875857
+    }
+  },
+  {
+    "epoch": 2,
+    "train_loss": 0.09120039182344977,
+    "tree_iou_old": 0.7348800379399489,
+    "tree_recall_new": 0.9992896559386689,
+    "miou_old_7": 0.7107307502748205,
+    "avg_tree": 0.8670848469393089,
+    "per_class_iou_old": {
+      "tree": 0.7348800379399489,
+      "ground": 0.853819505788582,
+      "person": 0.7223546582340064,
+      "sky": 0.7675516469276799,
+      "road": 0.7742544234364183,
+      "mountain": 0.41988628752707635,
+      "building": 0.7023686920700309,
+      "background": 0.32039096359921654
+    }
+  },
+  {
+    "epoch": 3,
+    "train_loss": 0.08554688886655733,
+    "tree_iou_old": 0.74153728943059,
+    "tree_recall_new": 0.9995758181795736,
+    "miou_old_7": 0.71243266621694,
+    "avg_tree": 0.8705565538050819,
+    "per_class_iou_old": {
+      "tree": 0.74153728943059,
+      "ground": 0.8535114953849786,
+      "person": 0.721128950704959,
+      "sky": 0.7701010375444943,
+      "road": 0.7727083269046633,
+      "mountain": 0.4298014657280298,
+      "building": 0.6982400978208655,
+      "background": 0.34584085232597234
+    }
+  },
+  {
+    "epoch": 4,
+    "train_loss": 0.08495914258924014,
+    "tree_iou_old": 0.7386737697019637,
+    "tree_recall_new": 0.9998306235571522,
+    "miou_old_7": 0.7108516336566473,
+    "avg_tree": 0.869252196629558,
+    "per_class_iou_old": {
+      "tree": 0.7386737697019637,
+      "ground": 0.8530840272060696,
+      "person": 0.7226071760588065,
+      "sky": 0.7702586411872955,
+      "road": 0.7661165021954849,
+      "mountain": 0.42538260855725873,
+      "building": 0.699838710689653,
+      "background": 0.3418799912840437
+    }
+  },
+  {
+    "epoch": 5,
+    "train_loss": 0.0824889050444397,
+    "tree_iou_old": 0.7374752361770358,
+    "tree_recall_new": 0.999875807068874,
+    "miou_old_7": 0.7075837051390604,
+    "avg_tree": 0.8686755216229549,
+    "per_class_iou_old": {
+      "tree": 0.7374752361770358,
+      "ground": 0.8517622487225316,
+      "person": 0.7170282863817369,
+      "sky": 0.7688366303765454,
+      "road": 0.7567630529761218,
+      "mountain": 0.42263572337088734,
+      "building": 0.6985847579685641,
+      "background": 0.3276835145430565
+    }
+  },
+  {
+    "epoch": 6,
+    "train_loss": NaN,
+    "tree_iou_old": 0.7397202992344901,
+    "tree_recall_new": 0.9999024727479229,
+    "miou_old_7": 0.7127320425514964,
+    "avg_tree": 0.8698113859912064,
+    "per_class_iou_old": {
+      "tree": 0.7397202992344901,
+      "ground": 0.853263524499737,
+      "person": 0.7202039547444198,
+      "sky": 0.7715238413193277,
+      "road": 0.7738255330134594,
+      "mountain": 0.4198599823783031,
+      "building": 0.7107271626707372,
+      "background": 0.34450614720925465
+    }
+  },
+  {
+    "epoch": 7,
+    "train_loss": NaN,
+    "tree_iou_old": 0.7398512261769766,
+    "tree_recall_new": 0.9999254348604373,
+    "miou_old_7": 0.7132055573889754,
+    "avg_tree": 0.869888330518707,
+    "per_class_iou_old": {
+      "tree": 0.7398512261769766,
+      "ground": 0.8539561450004414,
+      "person": 0.7196386093776438,
+      "sky": 0.7709864078362132,
+      "road": 0.7727917702136876,
+      "mountain": 0.4276638710715956,
+      "building": 0.707550872046269,
+      "background": 0.36019225517621223
+    }
+  },
+  {
+    "epoch": 8,
+    "train_loss": 0.07825338785261966,
+    "tree_iou_old": 0.7386024885258515,
+    "tree_recall_new": 0.9999464217374665,
+    "miou_old_7": 0.7114400608211693,
+    "avg_tree": 0.869274455131659,
+    "per_class_iou_old": {
+      "tree": 0.7386024885258515,
+      "ground": 0.8543745147138091,
+      "person": 0.72068353262178,
+      "sky": 0.7693893529932668,
+      "road": 0.7702198236853051,
+      "mountain": 0.4261176724920308,
+      "building": 0.700693040716142,
+      "background": 0.34337551501723146
+    }
+  }
+]

predict.py

@@ -21,7 +21,7 @@ from transformers import SegformerForSemanticSegmentation
 
 # ─── config ───
 BASE_MODEL = "nvidia/segformer-b5-finetuned-ade-640-640"
-WEIGHTS_PATH = os.environ.get("SEGFORMER85MV1_WEIGHTS", "Segformer85Mv1.pt")  # local file or full path
+WEIGHTS_PATH = os.environ.get("SEGFORMER85M_WEIGHTS", "Segformer85Mv2.pt")  # default v2 (best generalization); set env var or --weights to use v1
 NAMES = ["tree", "ground", "person", "sky", "road", "mountain", "building", "background"]
 PALETTE = np.array([
     [60, 220, 60],    # tree     - green