Modelpublic

Dnyanesh29/segformer-b2-desert

sourceHugging Faceapache-2.0updated 6mo agoView on Hugging Face

0likes4downloads

Model Card

🌵 SegFormer-B2 — Offroad Desert Semantic Segmentation

YOLO Pune Hackathon 2026 · Duality AI × MIT WPU

Fine-tuned nvidia/mit-b2 on synthetic desert imagery from Duality AI's Falcon simulation platform. The model segments every pixel of a desert scene into one of 10 classes. It was trained on Desert A and evaluated on a completely unseen Desert B location — a domain shift challenge.

📊 Results

Split	mIoU	Notes
Val (Desert A holdout)	0.6293	317 images, dedicated Duality AI split
Test (Desert B — raw)	0.2843	1002 images, different location
Test (Desert B — corrected)	0.4061	7 present classes only; Flowers, Logs, Ground Clutter absent in Desert B

Per-class IoU (Test Set — Desert B)

Class	Val IoU	Test IoU	Present in Desert B
🌳 Trees	0.8572	0.3896	✅
🌿 Lush Bushes	0.6937	0.0003	✅
🌾 Dry Grass	0.6923	0.4502	✅
🪨 Dry Bushes	0.4929	0.3831	✅
🌸 Flowers	0.5757	0.0000	❌ Absent
🪵 Logs	0.5214	0.0000	❌ Absent
⛰️ Rocks	0.4877	0.0402	✅
🏜️ Landscape	0.6007	0.5993	✅
☁️ Sky	0.9839	0.9802	✅
🪨 Ground Clutter	0.3874	0.0000	❌ Absent

Note on corrected mIoU: The 3 absent classes (Flowers, Logs, Ground Clutter) score IoU=0 by definition — the model never sees them in the test set. Corrected mIoU averages only the 7 classes that actually appear in Desert B.

🏗️ Model Details

Property	Value
Architecture	SegFormer-B2 (Mix-Transformer encoder + All-MLP decoder)
Base model	nvidia/mit-b2 (ImageNet-1K pretrained)
Total parameters	27.4M
Encoder	23.7M (pretrained)
Decoder head	3.7M (randomly initialised, fine-tuned)
Input resolution	512 × 512 px
Output classes	10
Training platform	Kaggle GPU (T4/P100)

🗂️ Dataset

Synthetic desert images generated by Duality AI's Falcon simulation platform.

Split	Images	Masks	Location
Train	2,857	2,857	Desert A
Val	317	317	Desert A (dedicated split)
Test	1,002	1,002	Desert B (unseen)

Mask label IDs — non-standard sparse integers remapped to 0–9:

Raw ID	Class	Compact ID
100	Trees	0
200	Lush Bushes	1
300	Dry Grass	2
500	Dry Bushes	3
550	Ground Clutter	4
600	Flowers ⚑ rare	5
700	Logs ⚑ rare	6
800	Rocks	7
7100	Landscape	8
10000	Sky	9

⚙️ Training Configuration

Optimiser & Schedule

Optimiser   : AdamW   lr=6e-5   weight_decay=1e-2
LR schedule : CosineAnnealingLR   T_max=50   eta_min=1e-7
Grad clip   : max_norm=1.0
Early stop  : patience=7 epochs on val mIoU
Best epoch  : 19   (V2 model)

Loss Function

Weighted CrossEntropyLoss with ignore_index=255 to handle unlabelled pixels.

Class	Weight	Reason
Flowers	5.0×	Extremely rare — forces model to notice them
Logs	4.0×	Rare and often partially occluded
Ground Clutter	2.0×	Small objects, easily missed
Lush Bushes, Dry Bushes	2.0×	Medium frequency
Trees, Dry Grass, Rocks	1.5×	Moderate
Sky	0.8×	Dominant — downweighted
Landscape	0.5×	Most dominant — heavily downweighted

Augmentation Pipeline

All four techniques recommended in Duality AI's official training guide:

python

A.Resize(512, 512),
A.HorizontalFlip(p=0.5),                         # 1. Flip — deserts have no L/R bias
A.RandomResizedCrop(size=(512,512),               # 2. Zoom — simulates camera distance
    scale=(0.6, 1.0), ratio=(0.75, 1.33), p=0.5),
A.Affine(shear=(-15,15), rotate=(-10,10), p=0.4), # 3. Shear — off-level camera angles
# 4. Mosaic — 4 images stitched into 2×2 grid (30% probability, custom implementation)
A.HueSaturationValue(                             # 5. HSV — MILD to preserve warm desert palette
    hue_shift_limit=10, sat_shift_limit=20,
    val_shift_limit=15, p=0.4),
A.RandomBrightnessContrast(
    brightness_limit=0.15, contrast_limit=0.15, p=0.4),
A.Normalize(mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]),

⚠️ HSV jitter is intentionally mild. The Duality AI desert palette is warm-toned. Aggressive colour shifts would train the model on unrealistic lighting and hurt generalisation.

🚀 Usage

Quick inference

python

import torch
import numpy as np
from PIL import Image
from transformers import SegformerForSemanticSegmentation
import albumentations as A
from albumentations.pytorch import ToTensorV2
import torch.nn.functional as F

# ── Class definitions ────────────────────────────────────────────────────────
CLASS_NAMES = [
    "Trees", "Lush_Bushes", "Dry_Grass", "Dry_Bushes",
    "Ground_Clutter", "Flowers", "Logs", "Rocks", "Landscape", "Sky"
]
PALETTE = np.array([
    [34,139,34],[0,200,83],[210,180,140],[200,200,180],[139,90,43],
    [255,20,147],[139,69,19],[128,128,128],[205,170,100],[135,206,235]
], dtype=np.uint8)

# ── Load model ───────────────────────────────────────────────────────────────
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = SegformerForSemanticSegmentation.from_pretrained(
    "YOUR_HF_USERNAME/segformer-b2-desert-segmentation"
).to(device).eval()

# ── Preprocess ───────────────────────────────────────────────────────────────
transform = A.Compose([
    A.Resize(512, 512),
    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ToTensorV2(),
])

# ── Inference ────────────────────────────────────────────────────────────────
img    = np.array(Image.open("desert_image.png").convert("RGB"))
h0, w0 = img.shape[:2]
tensor = transform(image=img)["image"].unsqueeze(0).to(device)

with torch.no_grad():
    logits = model(pixel_values=tensor).logits          # [1, 10, H/4, W/4]
    up     = F.interpolate(logits, size=(h0, w0),
                           mode="bilinear", align_corners=False)
    pred   = up.argmax(dim=1).squeeze(0).cpu().numpy()  # [H, W]  values 0–9

# ── Colour visualisation ─────────────────────────────────────────────────────
def mask_to_rgb(mask):
    rgb = np.zeros((*mask.shape, 3), dtype=np.uint8)
    for c in range(10):
        rgb[mask == c] = PALETTE[c]
    return rgb

pred_rgb = mask_to_rgb(pred)
Image.fromarray(pred_rgb).save("segmentation_output.png")
print("Classes found:", [CLASS_NAMES[c] for c in np.unique(pred)])

Load from checkpoint

python

import torch
from transformers import SegformerForSemanticSegmentation

# HuggingFace directory
model = SegformerForSemanticSegmentation.from_pretrained(
    "YOUR_HF_USERNAME/segformer-b2-desert-segmentation"
).eval()

# Raw .pth checkpoint (if you downloaded it separately)
ckpt = torch.load("best_model.pth", map_location="cpu")
model.load_state_dict(ckpt["model_state"])
print(f"Loaded epoch {ckpt['epoch']} — val mIoU {ckpt['val_miou']:.4f}")

⚠️ Limitations & Known Issues

—Domain shift is real. The model was trained on Desert A (Duality AI Falcon synthetic data). Performance on real-world desert images or different Falcon biomes may vary significantly.
—3 classes absent in Desert B. Flowers, Logs, and Ground Clutter do not appear in the test location. The model has learned to predict them on Desert A but will produce near-zero IoU on any location where they are absent.
—Lush Bushes test generalisation is poor (IoU 0.6937 val → 0.0003 test). Desert B appears to have a very different bush distribution or colour tone from Desert A.
—Rocks generalise weakly (0.4877 → 0.0402). Rock textures vary heavily between locations.
—Sky and Landscape generalise near-perfectly (Sky: 0.9839 → 0.9802; Landscape: 0.6007 → 0.5993). These classes are visually consistent across desert biomes.

📁 Repository Structure

segformer-b2-desert-segmentation/
├── config.json                  ← model architecture config
├── model.safetensors            ← fine-tuned weights (404 MB)
├── preprocessor_config.json     ← image processor settings
├── metadata.json                ← training metadata & scores
└── README.md                    ← this file

📚 Citation

If you use this model, please cite:

bibtex

@misc{mitwpu2025desert,
  title        = {SegFormer-B2 Fine-tuned on Duality AI Desert Segmentation},
  author       = {MIT WPU Team},
  year         = {2025},
  howpublished = {YOLO Pune Hackathon 2025, Duality AI Challenge},
  url          = {https://huggingface.co/YOUR_HF_USERNAME/segformer-b2-desert-segmentation}
}

Base model:

bibtex

@article{xie2021segformer,
  title   = {SegFormer: Simple and Efficient Design for Semantic Segmentation with Transformers},
  author  = {Xie, Enze and Wang, Wenhai and Yu, Zhiding and Anandkumar, Anima and Alvarez, Jose M and Luo, Ping},
  journal = {NeurIPS},
  year    = {2021}
}

🏅 Acknowledgements

—Duality AI for the Falcon synthetic dataset and challenge
—NVIDIA for the pretrained SegFormer-B2 backbone
—YOLO Pune Hackathon 2025 organisers at MIT WPU

Model trained and evaluated by MIT WPU for the Duality AI Offroad Segmentation challenge at YOLO Pune Hackathon 2026.