scripts/dsifn_proxy_train.py

"""Train a proxy DSIFN checkpoint (official weights unavailable on Drive).

Official Google Drive package only ships train/val/test zips — no .h5/.pth.
We therefore train the official PyTorch DSIFN (VGG16 ImageNet backbone)
on the DSIFN-CD val split for a few epochs and evaluate on the DSIFN test
split, then reuse the checkpoint for Delhi boundary comparison.

Usage:
    python scripts/dsifn_proxy_train.py --epochs 5 --batch 4
"""
from __future__ import annotations

import argparse
import importlib.util
import json
import sys
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from PIL import Image
from torch.utils.data import DataLoader, Dataset

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

DSIFN_PY = ROOT / "third_party" / "DSIFN" / "pytorch version" / "DSIFN.py"
DATA = ROOT / "third_party" / "DSIFN_weights" / "data"
OUT = ROOT / "models" / "dsifn_proxy"


def _load_dsifn_module():
    spec = importlib.util.spec_from_file_location("dsifn_official", DSIFN_PY)
    mod = importlib.util.module_from_spec(spec)
    assert spec.loader is not None
    spec.loader.exec_module(mod)
    return mod


class DSIFNPairDataset(Dataset):
    def __init__(self, split_dir: Path, size: int = 512):
        self.size = size
        self.t1_dir = split_dir / "t1"
        self.t2_dir = split_dir / "t2"
        self.mask_dir = split_dir / "mask"
        ids = []
        for p in sorted(self.t1_dir.iterdir()):
            stem = p.stem
            # mask may be .tif while images are .jpg
            m_candidates = list(self.mask_dir.glob(f"{stem}.*"))
            t2_candidates = list(self.t2_dir.glob(f"{stem}.*"))
            if m_candidates and t2_candidates:
                ids.append((p, t2_candidates[0], m_candidates[0], stem))
        self.items = ids

    def __len__(self):
        return len(self.items)

    def __getitem__(self, idx):
        t1p, t2p, mp, stem = self.items[idx]
        t1 = Image.open(t1p).convert("RGB").resize((self.size, self.size), Image.BILINEAR)
        t2 = Image.open(t2p).convert("RGB").resize((self.size, self.size), Image.BILINEAR)
        m = Image.open(mp).convert("L").resize((self.size, self.size), Image.NEAREST)
        t1 = torch.from_numpy(np.asarray(t1).transpose(2, 0, 1)).float() / 255.0
        t2 = torch.from_numpy(np.asarray(t2).transpose(2, 0, 1)).float() / 255.0
        # ImageNet normalize (VGG)
        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
        t1 = (t1 - mean) / std
        t2 = (t2 - mean) / std
        mask = (torch.from_numpy(np.array(m)).float() > 0).float().unsqueeze(0)
        return t1, t2, mask, stem


def _bce_multi(preds, target):
    # Deep supervision: average BCE over all branch outputs (resized to target)
    loss = 0.0
    for p in preds:
        if p.shape[-2:] != target.shape[-2:]:
            p = torch.nn.functional.interpolate(p, size=target.shape[-2:], mode="bilinear", align_corners=False)
        loss = loss + nn.functional.binary_cross_entropy(p.clamp(1e-6, 1 - 1e-6), target)
    return loss / float(len(preds))


@torch.no_grad()
def _eval_f1(model, loader, device, thr=0.5):
    model.eval()
    f1s = []
    for t1, t2, mask, _ in loader:
        t1, t2, mask = t1.to(device), t2.to(device), mask.to(device)
        outs = model(t1, t2)
        pred = outs[0]  # finest branch
        if pred.shape[-2:] != mask.shape[-2:]:
            pred = torch.nn.functional.interpolate(pred, size=mask.shape[-2:], mode="bilinear", align_corners=False)
        m = (pred >= thr).float()
        for i in range(m.shape[0]):
            gt = mask[i, 0] > 0.5
            pm = m[i, 0] > 0.5
            # Also try thr=0.3 if nothing fires at 0.5 (early training)
            if not gt.any():
                continue
            tp = (pm & gt).sum().item()
            fp = (pm & ~gt).sum().item()
            fn = (~pm & gt).sum().item()
            p = 0.0 if tp + fp == 0 else tp / (tp + fp)
            r = 0.0 if tp + fn == 0 else tp / (tp + fn)
            f1 = 0.0 if p + r == 0 else 2 * p * r / (p + r)
            f1s.append(f1)
    return float(np.mean(f1s)) if f1s else 0.0


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--epochs", type=int, default=5)
    ap.add_argument("--batch", type=int, default=2)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--size", type=int, default=256)  # memory-friendly; paper used 512
    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
    args = ap.parse_args()

    device = torch.device(args.device)
    OUT.mkdir(parents=True, exist_ok=True)

    mod = _load_dsifn_module()
    # Patch deprecated torchvision API if needed
    try:
        from torchvision.models import VGG16_Weights
        def _vgg():
            features = list(__import__("torchvision").models.vgg16(weights=VGG16_Weights.DEFAULT).features)[:30]
            m = torch.nn.Module()
            # rebuild like official
            base = mod.vgg16_base.__new__(mod.vgg16_base)
            torch.nn.Module.__init__(base)
            base.features = torch.nn.ModuleList(features).eval()
            return base
        model_a, model_b = _vgg(), _vgg()
    except Exception:
        model_a, model_b = mod.vgg16_base(), mod.vgg16_base()

    model = mod.DSIFN(model_a, model_b).to(device)
    # Freeze VGG backbone initially for faster proxy train
    for p in model.t1_base.parameters():
        p.requires_grad = False
    for p in model.t2_base.parameters():
        p.requires_grad = False

    train_ds = DSIFNPairDataset(DATA / "val", size=args.size)
    test_ds = DSIFNPairDataset(DATA / "test", size=args.size)
    print(f"train(val-split)={len(train_ds)} test={len(test_ds)} device={device}", flush=True)
    train_loader = DataLoader(train_ds, batch_size=args.batch, shuffle=True, num_workers=0)
    test_loader = DataLoader(test_ds, batch_size=args.batch, shuffle=False, num_workers=0)

    opt = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=args.lr)
    history = []
    best_f1, best_path = -1.0, OUT / "best.pt"

    for epoch in range(1, args.epochs + 1):
        model.train()
        # keep VGG in eval (BN frozen behavior)
        model.t1_base.eval()
        model.t2_base.eval()
        losses = []
        t0 = time.perf_counter()
        for t1, t2, mask, _ in train_loader:
            t1, t2, mask = t1.to(device), t2.to(device), mask.to(device)
            opt.zero_grad()
            outs = model(t1, t2)
            loss = _bce_multi(outs, mask)
            loss.backward()
            opt.step()
            losses.append(float(loss.item()))
        train_loss = float(np.mean(losses)) if losses else 0.0
        test_f1 = _eval_f1(model, test_loader, device)
        row = {"epoch": epoch, "train_loss": round(train_loss, 4), "test_f1": round(test_f1, 4),
               "elapsed_s": round(time.perf_counter() - t0, 1)}
        history.append(row)
        print(f"epoch {epoch}: loss={train_loss:.4f} test_f1={test_f1:.4f} ({row['elapsed_s']}s)", flush=True)
        if test_f1 > best_f1:
            best_f1 = test_f1
            torch.save({"model": model.state_dict(), "epoch": epoch, "test_f1": test_f1, "size": args.size}, best_path)

    meta = {
        "note": "Proxy DSIFN: official pretrained weights missing from Drive; trained on val, eval on test.",
        "epochs": args.epochs,
        "batch": args.batch,
        "lr": args.lr,
        "size": args.size,
        "best_test_f1": best_f1,
        "best_path": str(best_path),
        "history": history,
        "source_repo": "https://github.com/GeoZcx/A-deeply-supervised-image-fusion-network-for-change-detection-in-remote-sensing-images",
        "review_index": "https://github.com/MinZHANG-WHU/Change-Detection-Review",
    }
    (OUT / "metrics.json").write_text(json.dumps(meta, indent=2), encoding="utf-8")
    print(f"Best test F1={best_f1:.4f} -> {best_path}", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())