"""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()