Upload 9 files

add source training code. Weights landing in `weights` are in the root of the folder.

train/dataset.py

@@ -0,0 +1,93 @@
+"""
+dataset.py — PyTorch datasets wrapping the BDAPPV HuggingFace dataset.
+
+Two dataset classes:
+  - SegmentationDataset : positives only (has_mask=True), returns (image, mask)
+  - ClassificationDataset : all examples, returns (image, label)
+"""
+
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms as T
+
+IMG_MEAN = [0.485, 0.456, 0.406]
+IMG_STD  = [0.229, 0.224, 0.225]
+
+
+class SegmentationDataset(Dataset):
+    """
+    Wraps a HF split pre-filtered to positives (has_mask=True).
+    Returns (image_tensor [3, H, W], mask_tensor [1, H, W]) as float32.
+    """
+
+    def __init__(self, hf_split, img_size: int = 400, augment: bool = False):
+        self.data    = hf_split.filter(lambda x: x["has_mask"])
+        self.augment = augment
+        self.img_size = img_size
+
+        self.img_tf = T.Compose([
+            T.Resize((img_size, img_size)),
+            T.ToTensor(),
+            T.Normalize(mean=IMG_MEAN, std=IMG_STD),
+        ])
+        self.mask_tf = T.Compose([
+            T.Resize((img_size, img_size), interpolation=T.InterpolationMode.NEAREST),
+            T.ToTensor(),
+        ])
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        ex    = self.data[idx]
+        image = ex["image"].convert("RGB")
+        mask  = ex["mask"].convert("L")
+
+        if self.augment:
+            if torch.rand(1) > 0.5:
+                image = T.functional.hflip(image)
+                mask  = T.functional.hflip(mask)
+            if torch.rand(1) > 0.5:
+                image = T.functional.vflip(image)
+                mask  = T.functional.vflip(mask)
+
+        image = self.img_tf(image)
+        mask  = self.mask_tf(mask)
+        mask  = (mask > 0.5).float()   # binarize → {0, 1}
+
+        return image, mask
+
+
+class ClassificationDataset(Dataset):
+    """
+    All examples (positives + negatives).
+    Returns (image_tensor [3, H, W], label [float scalar]).
+    label = 1.0 if has_mask else 0.0
+    """
+
+    def __init__(self, hf_split, img_size: int = 299, augment: bool = False):
+        self.data    = hf_split
+        self.augment = augment
+
+        self.img_tf = T.Compose([
+            T.Resize((img_size, img_size)),
+            T.ToTensor(),
+            T.Normalize(mean=IMG_MEAN, std=IMG_STD),
+        ])
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        ex    = self.data[idx]
+        image = ex["image"].convert("RGB")
+        label = float(ex["has_mask"])
+
+        if self.augment:
+            if torch.rand(1) > 0.5:
+                image = T.functional.hflip(image)
+            if torch.rand(1) > 0.5:
+                image = T.functional.vflip(image)
+
+        image = self.img_tf(image)
+        return image, torch.tensor(label)

train/evaluate.py

@@ -0,0 +1,162 @@
+"""
+evaluate.py — Evaluate trained models on the BDAPPV test split.
+
+Usage:
+    # Segmentation
+    python evaluate.py seg --provider google --checkpoint weights/deeplab_google_best.pth
+    python evaluate.py seg --provider ign    --checkpoint weights/deeplab_ign_best.pth
+
+    # Classification
+    python evaluate.py clf --provider google --checkpoint weights/inception_google_best.pth
+
+    # Cross-provider (train Google, eval IGN) — the distribution shift benchmark
+    python evaluate.py seg --provider ign --checkpoint weights/deeplab_google_best.pth
+"""
+
+import argparse
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchvision.models.segmentation import deeplabv3_resnet101
+from torchvision.models import inception_v3
+
+from datasets import load_dataset
+
+from dataset import SegmentationDataset, ClassificationDataset
+
+
+# ── Metrics ───────────────────────────────────────────────────────────────────
+
+def compute_seg_metrics(pred_logits, target, threshold=0.5):
+    pred  = (torch.sigmoid(pred_logits) > threshold).float()
+    tp    = (pred * target).sum().item()
+    fp    = (pred * (1 - target)).sum().item()
+    fn    = ((1 - pred) * target).sum().item()
+    iou   = tp / max(tp + fp + fn, 1e-6)
+    prec  = tp / max(tp + fp, 1e-6)
+    rec   = tp / max(tp + fn, 1e-6)
+    f1    = 2 * prec * rec / max(prec + rec, 1e-6)
+    return iou, f1
+
+
+def compute_clf_metrics(logits, labels, threshold=0.0):
+    preds = (logits > threshold).float()
+    acc   = (preds == labels).float().mean().item()
+    tp    = (preds * labels).sum().item()
+    fp    = (preds * (1 - labels)).sum().item()
+    fn    = ((1 - preds) * labels).sum().item()
+    prec  = tp / max(tp + fp, 1e-6)
+    rec   = tp / max(tp + fn, 1e-6)
+    f1    = 2 * prec * rec / max(prec + rec, 1e-6)
+    return acc, prec, rec, f1
+
+
+# ── Segmentation eval ─────────────────────────────────────────────────────────
+
+def eval_seg(args):
+    device = (
+        torch.device("mps")  if torch.backends.mps.is_available() else
+        torch.device("cuda") if torch.cuda.is_available() else
+        torch.device("cpu")
+    )
+
+    ds      = load_dataset("gabrielkasmi/bdappv", args.provider)
+    test_ds = SegmentationDataset(ds["test"], img_size=args.img_size, augment=False)
+    loader  = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
+
+    model = deeplabv3_resnet101(weights=None, aux_loss=False)
+    model.classifier[-1] = nn.Conv2d(256, 1, kernel_size=1)
+    state      = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+    model_dict = model.state_dict()
+    compatible = {k: v for k, v in state.items() if k in model_dict and v.shape == model_dict[k].shape}
+    model_dict.update(compatible)
+    model.load_state_dict(model_dict)
+    model = model.to(device)
+    model.eval()
+
+    total_iou, total_f1 = 0.0, 0.0
+    with torch.no_grad():
+        for images, masks in loader:
+            images, masks = images.to(device), masks.to(device)
+            logits = model(images)["out"]
+            if logits.shape[-2:] != masks.shape[-2:]:
+                logits = nn.functional.interpolate(logits, size=masks.shape[-2:], mode="bilinear", align_corners=False)
+            iou, f1 = compute_seg_metrics(logits, masks)
+            total_iou += iou
+            total_f1  += f1
+
+    n = len(loader)
+    print(f"\n── Segmentation results ──────────────────────")
+    print(f"  Provider   : {args.provider}")
+    print(f"  Checkpoint : {args.checkpoint}")
+    print(f"  Test images: {len(test_ds):,}")
+    print(f"  IoU  : {total_iou / n:.4f}")
+    print(f"  F1   : {total_f1  / n:.4f}")
+
+
+# ── Classification eval ───────────────────────────────────────────────────────
+
+def eval_clf(args):
+    device = (
+        torch.device("mps")  if torch.backends.mps.is_available() else
+        torch.device("cuda") if torch.cuda.is_available() else
+        torch.device("cpu")
+    )
+
+    ds      = load_dataset("gabrielkasmi/bdappv", args.provider)
+    test_ds = ClassificationDataset(ds["test"], img_size=299, augment=False)
+    loader  = DataLoader(test_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers)
+
+    model              = inception_v3(weights=None, aux_logits=True)
+    model.fc           = nn.Linear(model.fc.in_features, 1)
+    model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, 1)
+    state              = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+    model.load_state_dict(state)
+    model.aux_logits   = False
+    model    = model.to(device)
+    model.eval()
+
+    all_logits, all_labels = [], []
+    with torch.no_grad():
+        for images, labels in loader:
+            images = images.to(device)
+            logits = model(images)
+            if hasattr(logits, "logits"):
+                logits = logits.logits
+            all_logits.append(logits.cpu())
+            all_labels.append(labels.unsqueeze(1))
+
+    all_logits = torch.cat(all_logits)
+    all_labels = torch.cat(all_labels)
+    acc, prec, rec, f1 = compute_clf_metrics(all_logits, all_labels)
+
+    print(f"\n── Classification results ────────────────────")
+    print(f"  Provider   : {args.provider}")
+    print(f"  Checkpoint : {args.checkpoint}")
+    print(f"  Test images: {len(test_ds):,}")
+    print(f"  Accuracy   : {acc:.4f}")
+    print(f"  Precision  : {prec:.4f}")
+    print(f"  Recall     : {rec:.4f}")
+    print(f"  F1         : {f1:.4f}")
+
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    sub    = parser.add_subparsers(dest="task", required=True)
+
+    for name in ["seg", "clf"]:
+        p = sub.add_parser(name)
+        p.add_argument("--provider",    required=True, choices=["google", "ign"])
+        p.add_argument("--checkpoint",  required=True)
+        p.add_argument("--batch_size",  type=int, default=8)
+        p.add_argument("--img_size",    type=int, default=400)
+        p.add_argument("--num_workers", type=int, default=0)
+
+    args = parser.parse_args()
+    if args.task == "seg":
+        eval_seg(args)
+    else:
+        eval_clf(args)

train/model.py

@@ -0,0 +1,74 @@
+"""
+model.py — helpers to load BDAPPV baseline models from HuggingFace Hub.
+
+Usage:
+    from huggingface_hub import hf_hub_download
+    import importlib.util, sys
+
+    path = hf_hub_download("gabrielkasmi/bdappv-models", "model.py")
+    spec = importlib.util.spec_from_file_location("bdappv_model", path)
+    mod  = importlib.util.module_from_spec(spec); spec.loader.exec_module(mod)
+
+    seg = mod.load_segmentation_model("google")   # or "ign"
+    clf = mod.load_classification_model("google") # or "ign"
+"""
+
+import torch
+import torch.nn as nn
+from huggingface_hub import hf_hub_download
+from torchvision.models.segmentation import deeplabv3_resnet101
+from torchvision.models import inception_v3
+
+REPO = "gabrielkasmi/bdappv-models"
+
+
+def load_segmentation_model(provider: str = "google", device: str = "cpu"):
+    """
+    Load the DeepLabV3-ResNet101 segmentation model.
+
+    Args:
+        provider : "google" or "ign"
+        device   : "cpu", "cuda", or "mps"
+
+    Returns:
+        model in eval mode
+    """
+    assert provider in ("google", "ign"), "provider must be 'google' or 'ign'"
+
+    path  = hf_hub_download(REPO, f"deeplab_{provider}_best.pth")
+    model = deeplabv3_resnet101(weights=None, aux_loss=False)
+    model.classifier[-1] = nn.Conv2d(256, 1, kernel_size=1)
+
+    state      = torch.load(path, map_location=device, weights_only=False)
+    model_dict = model.state_dict()
+    compatible = {k: v for k, v in state.items()
+                  if k in model_dict and v.shape == model_dict[k].shape}
+    model_dict.update(compatible)
+    model.load_state_dict(model_dict)
+
+    return model.eval().to(device)
+
+
+def load_classification_model(provider: str = "google", device: str = "cpu"):
+    """
+    Load the InceptionV3 classification model (panel / no panel).
+
+    Args:
+        provider : "google" or "ign"
+        device   : "cpu", "cuda", or "mps"
+
+    Returns:
+        model in eval mode
+    """
+    assert provider in ("google", "ign"), "provider must be 'google' or 'ign'"
+
+    path  = hf_hub_download(REPO, f"inception_{provider}_best.pth")
+    model = inception_v3(weights=None, aux_logits=True)
+    model.fc           = nn.Linear(model.fc.in_features, 1)
+    model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, 1)
+
+    state = torch.load(path, map_location=device, weights_only=False)
+    model.load_state_dict(state)
+    model.aux_logits = False  # désactive pour l'inférence
+
+    return model.eval().to(device)

train/pretrained/deeplabv3_weights.tar

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

train/pretrained/inceptionv3_weights.tar

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

train/requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.2.0
+torchvision>=0.17.0
+datasets>=2.18.0
+Pillow>=10.0.0

train/train_clf.py

@@ -0,0 +1,168 @@
+"""
+train_clf.py — InceptionV3 binary classification on BDAPPV (panel / no panel).
+
+Usage:
+    python train_clf.py --provider google
+    python train_clf.py --provider ign
+    python train_clf.py --provider google --checkpoint weights/inception_google_best.pth
+"""
+
+import argparse
+import os
+
+import torch
+from tqdm import tqdm
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchvision.models import inception_v3, Inception_V3_Weights
+
+from datasets import load_dataset
+
+from dataset import ClassificationDataset
+
+
+# ── Metrics ───────────────────────────────────────────────────────────────────
+
+def accuracy_f1(logits, labels, threshold=0.0):
+    preds   = (logits > threshold).float()
+    correct = (preds == labels).float().mean().item()
+    tp = (preds * labels).sum().item()
+    fp = (preds * (1 - labels)).sum().item()
+    fn = ((1 - preds) * labels).sum().item()
+    prec = tp / max(tp + fp, 1e-6)
+    rec  = tp / max(tp + fn, 1e-6)
+    f1   = 2 * prec * rec / max(prec + rec, 1e-6)
+    return correct, f1
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+
+def train(args):
+    device = (
+        torch.device("mps")  if torch.backends.mps.is_available() else
+        torch.device("cuda") if torch.cuda.is_available() else
+        torch.device("cpu")
+    )
+    print(f"Device : {device}")
+    print(f"Provider: {args.provider}  |  effective batch: {args.batch_size * args.grad_accum}")
+
+    # ── Data ──────────────────────────────────────────────────────────────────
+    print("Loading dataset from HuggingFace Hub...")
+    ds = load_dataset("gabrielkasmi/bdappv", args.provider)
+
+    train_ds = ClassificationDataset(ds["train"],      img_size=299, augment=True)
+    val_ds   = ClassificationDataset(ds["validation"], img_size=299, augment=False)
+
+    print(f"  train : {len(train_ds):,}  |  val : {len(val_ds):,}")
+
+    train_loader = DataLoader(
+        train_ds, batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, pin_memory=(device.type != "mps"),
+    )
+    val_loader = DataLoader(
+        val_ds, batch_size=args.batch_size, shuffle=False,
+        num_workers=args.num_workers, pin_memory=(device.type != "mps"),
+    )
+
+    # ── Model ─────────────────────────────────────────────────────────────────
+    model = inception_v3(weights=Inception_V3_Weights.DEFAULT)
+
+    if args.checkpoint:
+        # Poids issus d'un modèle à 2 classes — on charge avec la tête originale
+        # puis on la remplace par une tête binaire (1 sortie, BCEWithLogitsLoss)
+        model.fc           = nn.Linear(model.fc.in_features, 2)
+        model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, 2)
+        ckpt  = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+        state = ckpt["model_state_dict"] if isinstance(ckpt, dict) else ckpt
+        model.load_state_dict(state, strict=True)
+        print(f"Loaded checkpoint: {args.checkpoint}")
+
+    # Remplace la tête par 1 sortie pour BCEWithLogitsLoss
+    model.fc           = nn.Linear(model.fc.in_features, 1)
+    model.AuxLogits.fc = nn.Linear(model.AuxLogits.fc.in_features, 1)
+
+    model = model.to(device)
+
+    # ── Optimizer + scheduler ─────────────────────────────────────────────────
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+    criterion = nn.BCEWithLogitsLoss()
+
+    # ── Training loop ─────────────────────────────────────────────────────────
+    os.makedirs(args.out_dir, exist_ok=True)
+    best_f1      = 0.0
+    patience_ctr = 0
+
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        total_loss = 0.0
+        optimizer.zero_grad()
+
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch}/{args.epochs}", leave=False)
+        for step, (images, labels) in enumerate(pbar):
+            images = images.to(device)
+            labels = labels.to(device).unsqueeze(1)
+
+            out    = model(images)
+            logits = out.logits if hasattr(out, "logits") else out[0]
+            aux    = out.aux_logits if hasattr(out, "aux_logits") else out[1]
+
+            loss = (criterion(logits, labels) + 0.4 * criterion(aux, labels)) / args.grad_accum
+            loss.backward()
+            total_loss += loss.item() * args.grad_accum
+            pbar.set_postfix(loss=f"{loss.item() * args.grad_accum:.4f}")
+
+            if (step + 1) % args.grad_accum == 0 or (step + 1) == len(train_loader):
+                optimizer.step()
+                optimizer.zero_grad()
+
+        # ── Validation ────────────────────────────────────────────────────────
+        model.eval()
+        all_logits, all_labels = [], []
+        with torch.no_grad():
+            for images, labels in val_loader:
+                images = images.to(device)
+                logits = model(images)
+                if hasattr(logits, "logits"):
+                    logits = logits.logits
+                all_logits.append(logits.cpu())
+                all_labels.append(labels.unsqueeze(1))
+
+        all_logits = torch.cat(all_logits)
+        all_labels = torch.cat(all_labels)
+        acc, f1    = accuracy_f1(all_logits, all_labels)
+        avg_loss   = total_loss / len(train_loader)
+
+        print(f"Epoch {epoch:3d}/{args.epochs}  loss={avg_loss:.4f}  val_acc={acc:.4f}  val_f1={f1:.4f}")
+
+        scheduler.step()
+
+        if f1 > best_f1:
+            best_f1      = f1
+            patience_ctr = 0
+            ckpt_path    = os.path.join(args.out_dir, f"inception_{args.provider}_best.pth")
+            torch.save(model.state_dict(), ckpt_path)
+            print(f"  → checkpoint saved (F1={best_f1:.4f})")
+        else:
+            patience_ctr += 1
+            if patience_ctr >= args.patience:
+                print(f"\nEarly stopping — no improvement for {args.patience} epochs.")
+                break
+
+    print(f"\nDone. Best val F1: {best_f1:.4f}")
+    print(f"Weights: {os.path.join(args.out_dir, f'inception_{args.provider}_best.pth')}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--provider",    default="google", choices=["google", "ign"])
+    parser.add_argument("--epochs",      type=int,   default=50)
+    parser.add_argument("--batch_size",  type=int,   default=8)
+    parser.add_argument("--grad_accum",  type=int,   default=4,   help="Gradient accumulation steps")
+    parser.add_argument("--lr",          type=float, default=1e-4)
+    parser.add_argument("--num_workers", type=int,   default=2,   help="Workers pour le dataloader")
+    parser.add_argument("--checkpoint",  default=None)
+    parser.add_argument("--out_dir",     default="weights")
+    parser.add_argument("--patience",    type=int,   default=7,   help="Early stopping patience (epochs)")
+    args = parser.parse_args()
+    train(args)

train/train_seg.py

@@ -0,0 +1,168 @@
+"""
+train_seg.py — DeepLabV3-ResNet50 segmentation on BDAPPV.
+
+Usage:
+    python train_seg.py --provider google
+    python train_seg.py --provider ign
+    python train_seg.py --provider google --checkpoint weights/deeplab_google_best.pth  # resume
+"""
+
+import argparse
+import os
+
+import torch
+from tqdm import tqdm
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from torchvision.models.segmentation import deeplabv3_resnet101, DeepLabV3_ResNet101_Weights
+
+from datasets import load_dataset
+
+from dataset import SegmentationDataset
+
+
+# ── Metrics ───────────────────────────────────────────────────────────────────
+
+def iou_score(pred_logits, target, threshold=0.5):
+    pred  = (torch.sigmoid(pred_logits) > threshold).float()
+    inter = (pred * target).sum()
+    union = pred.sum() + target.sum() - inter
+    return (inter / union.clamp(min=1e-6)).item()
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+
+def train(args):
+    device = (
+        torch.device("mps")  if torch.backends.mps.is_available() else
+        torch.device("cuda") if torch.cuda.is_available() else
+        torch.device("cpu")
+    )
+    print(f"Device : {device}")
+    print(f"Provider: {args.provider}  |  effective batch: {args.batch_size * args.grad_accum}")
+
+    # ── Data ──────────────────────────────────────────────────────────────────
+    print("Loading dataset from HuggingFace Hub...")
+    ds = load_dataset("gabrielkasmi/bdappv", args.provider)
+
+    train_ds = SegmentationDataset(ds["train"],      img_size=args.img_size, augment=True)
+    val_ds   = SegmentationDataset(ds["validation"], img_size=args.img_size, augment=False)
+
+    print(f"  train positives : {len(train_ds):,}")
+    print(f"  val   positives : {len(val_ds):,}")
+
+    train_loader = DataLoader(
+        train_ds, batch_size=args.batch_size, shuffle=True,
+        num_workers=args.num_workers, pin_memory=(device.type != "mps"),
+        drop_last=True,
+    )
+    val_loader = DataLoader(
+        val_ds, batch_size=args.batch_size, shuffle=False,
+        num_workers=args.num_workers, pin_memory=(device.type != "mps"),
+    )
+
+    # ── Model ─────────────────────────────────────────────────────────────────
+    model = deeplabv3_resnet101(weights=DeepLabV3_ResNet101_Weights.DEFAULT)
+    model.classifier[-1]     = nn.Conv2d(256, 1, kernel_size=1)
+    model.aux_classifier[-1] = nn.Conv2d(256, 1, kernel_size=1)
+
+    if args.checkpoint:
+        # Le checkpoint mélange une tête principale binaire et une tête aux à 21 classes
+        # (poids de base incohérents) — on ne charge que les tenseurs dont la shape colle.
+        ckpt  = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+        state = ckpt["model_state_dict"] if isinstance(ckpt, dict) else ckpt
+        model_dict = model.state_dict()
+        compatible = {k: v for k, v in state.items() if k in model_dict and v.shape == model_dict[k].shape}
+        skipped = set(state.keys()) - set(compatible.keys())
+        model_dict.update(compatible)
+        model.load_state_dict(model_dict)
+        print(f"Loaded checkpoint: {args.checkpoint}  ({len(compatible)}/{len(state)} tensors matched, {len(skipped)} skipped)")
+
+    model = model.to(device)
+
+    # ── Optimizer + scheduler ─────────────────────────────────────────────────
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+    criterion = nn.BCEWithLogitsLoss()
+
+    # ── Training loop ─────────────────────────────────────────────────────────
+    os.makedirs(args.out_dir, exist_ok=True)
+    best_iou     = 0.0
+    patience_ctr = 0
+
+    for epoch in range(1, args.epochs + 1):
+        model.train()
+        total_loss = 0.0
+        optimizer.zero_grad()
+
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch}/{args.epochs}", leave=False)
+        for step, (images, masks) in enumerate(pbar):
+            if step == 3:
+                break
+            images, masks = images.to(device), masks.to(device)
+
+            out    = model(images)
+            logits = out["out"]
+            aux    = out["aux"]
+
+            if logits.shape[-2:] != masks.shape[-2:]:
+                logits = nn.functional.interpolate(logits, size=masks.shape[-2:], mode="bilinear", align_corners=False)
+                aux    = nn.functional.interpolate(aux,    size=masks.shape[-2:], mode="bilinear", align_corners=False)
+
+            loss = (criterion(logits, masks) + 0.4 * criterion(aux, masks)) / args.grad_accum
+            loss.backward()
+            total_loss += loss.item() * args.grad_accum
+            pbar.set_postfix(loss=f"{loss.item() * args.grad_accum:.4f}")
+
+            if (step + 1) % args.grad_accum == 0 or (step + 1) == len(train_loader):
+                optimizer.step()
+                optimizer.zero_grad()
+
+        # ── Validation ────────────────────────────────────────────────────────
+        model.eval()
+        val_iou = 0.0
+        with torch.no_grad():
+            for images, masks in val_loader:
+                images, masks = images.to(device), masks.to(device)
+                logits = model(images)["out"]
+                if logits.shape[-2:] != masks.shape[-2:]:
+                    logits = nn.functional.interpolate(logits, size=masks.shape[-2:], mode="bilinear", align_corners=False)
+                val_iou += iou_score(logits, masks)
+
+        val_iou   /= len(val_loader)
+        avg_loss   = total_loss / len(train_loader)
+
+        print(f"Epoch {epoch:3d}/{args.epochs}  loss={avg_loss:.4f}  val_iou={val_iou:.4f}")
+
+        scheduler.step()
+
+        if val_iou > best_iou:
+            best_iou     = val_iou
+            patience_ctr = 0
+            ckpt_path    = os.path.join(args.out_dir, f"deeplab_{args.provider}_best.pth")
+            torch.save(model.state_dict(), ckpt_path)
+            print(f"  → checkpoint saved (IoU={best_iou:.4f})")
+        else:
+            patience_ctr += 1
+            if patience_ctr >= args.patience:
+                print(f"\nEarly stopping — no improvement for {args.patience} epochs.")
+                break
+
+    print(f"\nDone. Best val IoU: {best_iou:.4f}")
+    print(f"Weights: {os.path.join(args.out_dir, f'deeplab_{args.provider}_best.pth')}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--provider",    default="google", choices=["google", "ign"])
+    parser.add_argument("--epochs",      type=int,   default=50)
+    parser.add_argument("--batch_size",  type=int,   default=4)
+    parser.add_argument("--grad_accum",  type=int,   default=8,   help="Gradient accumulation steps (eff. batch = batch_size × grad_accum)")
+    parser.add_argument("--lr",          type=float, default=1e-4)
+    parser.add_argument("--img_size",    type=int,   default=400)
+    parser.add_argument("--num_workers", type=int,   default=2,   help="Workers pour le dataloader")
+    parser.add_argument("--checkpoint",  default=None)
+    parser.add_argument("--out_dir",     default="weights")
+    parser.add_argument("--patience",    type=int,   default=7,   help="Early stopping patience (epochs)")
+    args = parser.parse_args()
+    train(args)