| |
| """ |
| setup.py - One-command setup for the CD-Models benchmark suite. |
| |
| Clones all model repositories, downloads pretrained backbone weights, |
| verifies the Python environment, and optionally prepares datasets. |
| |
| Usage: |
| python setup.py # full setup (clone + weights + env check) |
| python setup.py --skip-weights # clone only, skip weight downloads |
| python setup.py --env-check-only # only check Python packages |
| python setup.py --status # print current setup status without changes |
| python setup.py --dataset levir_cd # prepare list files for a specific dataset |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import importlib |
| import os |
| import subprocess |
| import sys |
| import urllib.request |
| from pathlib import Path |
|
|
|
|
| ROOT = Path(__file__).resolve().parent |
| DATASET_CONFIGS = ROOT / "configs" / "datasets" |
| MODEL_REPOS = ROOT / "model_repos" |
| IMG_EXTS = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"} |
| ROOT_LEVEL_REPOS = { |
| "BIT_CD", |
| "ChangeFormer", |
| "STANet", |
| "Change3D", |
| "BiFA", |
| "CDMamba", |
| "RSM-CD", |
| "SChanger", |
| "IFNet", |
| "Siam-NestedUNet", |
| } |
|
|
| REQUIRED_PACKAGES = [ |
| ("torch", "torch", "all models", False), |
| ("torchvision", "torchvision", "all models", False), |
| ("numpy", "numpy", "all models", True), |
| ("PIL", "Pillow", "all models", True), |
| ("cv2", "opencv-python", "all models", True), |
| ("yaml", "pyyaml", "all models", True), |
| ("tqdm", "tqdm", "weight downloads", True), |
| ("timm", "timm", "TinyCD, ELGC-Net, ChangeFormer", True), |
| ("einops", "einops", "BIT_CD, ChangeFormer", True), |
| ("sklearn", "scikit-learn", "metrics", True), |
| ("gdown", "gdown", "GDrive weight fallback", True), |
| ("mmengine", "mmengine==0.10.1", "ChangeMamba, Changer", False), |
| ("mmcv", "mmcv==2.1.0", "ChangeMamba, Changer", False), |
| ("mmseg", "mmsegmentation==1.2.2", "ChangeMamba, Changer", False), |
| ] |
|
|
| CONFIRMED_REPOS = { |
| "BIT_CD": "https://github.com/justchenhao/BIT_CD.git", |
| "ChangeFormer": "https://github.com/wgcban/ChangeFormer.git", |
| "STANet": "https://github.com/justchenhao/STANet.git", |
| "Change3D": "https://github.com/Z-Zheng/Change3D.git", |
| "BiFA": "https://github.com/zmoka-zht/BiFA.git", |
| "CDMamba": "https://github.com/zmoka-zht/CDMamba.git", |
| "RSM-CD": "https://github.com/walking-shadow/Official_Remote_Sensing_Mamba.git", |
| "SChanger": None, |
| "IFNet": "https://github.com/GeoZcx/A-deeply-supervised-image-fusion-network-for-change-detection.git", |
| "Siam-NestedUNet": "https://github.com/likyoo/Siam-NestedUNet.git", |
| "fully_convolutional_change_detection": "https://github.com/rcdaudt/fully_convolutional_change_detection.git", |
| "ChangeMamba": "https://github.com/ChenHongruixuan/ChangeMamba.git", |
| "DSAMNet": "https://github.com/liumency/DSAMNet.git", |
| "Tiny_model_4_CD": "https://github.com/AndreaCodegoni/Tiny_model_4_CD.git", |
| "HANet-CD": "https://github.com/ChengxiHAN/HANet-CD.git", |
| "CGNet-CD": "https://github.com/ChengxiHAN/CGNet-CD.git", |
| "open-cd": "https://github.com/likyoo/open-cd.git", |
| "elgcnet": "https://github.com/techmn/elgcnet.git", |
| } |
|
|
| ZENODO_WEIGHTS = { |
| "vmamba_tiny": { |
| "record_id": "14037770", |
| "filename": "vssmtiny_dp01_ckpt_epoch_292.pth", |
| "dest_dir": MODEL_REPOS / "ChangeMamba" / "pretrained_weight", |
| "sha256": None, |
| "notes": "VMamba-Tiny for ChangeMamba", |
| }, |
| "vmamba_small": { |
| "record_id": "14037770", |
| "filename": "vssmsmall_dp03_ckpt_epoch_238.pth", |
| "dest_dir": MODEL_REPOS / "ChangeMamba" / "pretrained_weight", |
| "sha256": None, |
| "notes": "VMamba-Small for ChangeMamba", |
| }, |
| "vmamba_base": { |
| "record_id": "14037770", |
| "filename": "vssmbase_dp06_ckpt_epoch_241.pth", |
| "dest_dir": MODEL_REPOS / "ChangeMamba" / "pretrained_weight", |
| "sha256": None, |
| "notes": "VMamba-Base for ChangeMamba", |
| }, |
| } |
|
|
| TIMM_WEIGHTS_TO_PREFETCH = [ |
| "efficientnet_b4", |
| "mit_b0", |
| "mit_b1", |
| "mit_b4", |
| ] |
|
|
| TORCHVISION_WEIGHTS_TO_PREFETCH = [ |
| ("resnet18", "IMAGENET1K_V1"), |
| ("resnet50", "IMAGENET1K_V1"), |
| ("vgg16", "IMAGENET1K_V1"), |
| ] |
|
|
| TIMM_ALIASES = { |
| "efficientnet_b4": ("efficientnet_b4", "tf_efficientnet_b4", "tf_efficientnet_b4_ns"), |
| "mit_b0": ("mit_b0", "segformer_b0"), |
| "mit_b1": ("mit_b1", "segformer_b1"), |
| "mit_b4": ("mit_b4", "segformer_b4"), |
| } |
|
|
|
|
| def run(cmd: list[str]) -> subprocess.CompletedProcess: |
| return subprocess.run(cmd, capture_output=True, text=True, check=False) |
|
|
|
|
| def import_status(import_name: str) -> tuple[bool, str]: |
| try: |
| module = importlib.import_module(import_name) |
| except Exception as exc: |
| return False, str(exc) |
| return True, getattr(module, "__version__", "present") |
|
|
|
|
| def pip_install(package: str) -> bool: |
| print(f" [INSTALL] pip install {package}") |
| result = run([sys.executable, "-m", "pip", "install", package]) |
| if result.returncode != 0: |
| print(f" [FAIL] {package}: {result.stderr[-500:]}") |
| return False |
| print(f" [OK] {package} installed") |
| return True |
|
|
|
|
| def check_environment(status_only: bool = False) -> dict[str, int]: |
| print("\nPYTHON ENVIRONMENT") |
| print("------------------") |
| print(f"Python executable: {sys.executable}") |
| print(f"Python version: {sys.version.split()[0]}") |
| present = missing = installed = warnings = 0 |
|
|
| for import_name, pip_name, required_for, auto_install in REQUIRED_PACKAGES: |
| ok, detail = import_status(import_name) |
| if ok: |
| present += 1 |
| print(f" [OK] {import_name:<12} {detail:<18} required for {required_for}") |
| continue |
| missing += 1 |
| if auto_install and not status_only: |
| if pip_install(pip_name): |
| installed += 1 |
| continue |
| if not auto_install: |
| warnings += 1 |
| print(f" [WARN] {import_name:<12} missing; install manually: pip install {pip_name}") |
| else: |
| print(f" [MISSING] {import_name:<12} install: pip install {pip_name}") |
|
|
| if warnings: |
| print("\nVersion-sensitive packages are intentionally not auto-installed.") |
| print("For ChangeMamba / Changer:") |
| print(" pip install mmengine==0.10.1") |
| print(" pip install mmcv==2.1.0 -f https://download.openmmlab.com/mmcv/dist/cu124/torch2.6/index.html") |
| print(" pip install mmsegmentation==1.2.2 mmdet==3.3.0 mmpretrain==1.2.0") |
|
|
| return {"present": present, "missing": missing, "installed": installed, "warnings": warnings} |
|
|
|
|
| def clone_repo(name: str, url: str | None, dest_dir: Path = MODEL_REPOS, status_only: bool = False) -> bool: |
| dest = (ROOT / name) if name in ROOT_LEVEL_REPOS else (dest_dir / name) |
| if dest.is_dir(): |
| print(f" [OK] {name} already cloned") |
| return True |
| if url is None: |
| print(f" [SKIP] {name} - URL not confirmed, clone manually") |
| return False |
| if status_only: |
| print(f" [MISSING] {name} -> {url}") |
| return False |
| print(f" [CLONE] {name}...") |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| result = run(["git", "clone", "--depth", "1", url, str(dest)]) |
| if result.returncode == 0: |
| print(f" [OK] {name} cloned") |
| return True |
| print(f" [FAIL] {name}: {result.stderr[:300]}") |
| return False |
|
|
|
|
| def clone_repositories(status_only: bool = False) -> dict[str, int]: |
| print("\nMODEL REPOSITORIES") |
| print("------------------") |
| ok = missing = skipped = 0 |
| for name, url in CONFIRMED_REPOS.items(): |
| result = clone_repo(name, url, status_only=status_only) |
| if result: |
| ok += 1 |
| elif url is None: |
| skipped += 1 |
| else: |
| missing += 1 |
| return {"ok": ok, "missing": missing, "skipped": skipped, "total": len(CONFIRMED_REPOS)} |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as f: |
| for chunk in iter(lambda: f.read(1024 * 1024), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def download_zenodo_weight( |
| record_id: str, |
| filename: str, |
| dest_dir: Path, |
| expected_sha256: str | None = None, |
| status_only: bool = False, |
| ) -> Path | None: |
| dest_dir.mkdir(parents=True, exist_ok=True) |
| dest_path = dest_dir / filename |
| if dest_path.exists(): |
| if expected_sha256: |
| actual = sha256_file(dest_path) |
| if actual != expected_sha256: |
| print(f" [FAIL] {filename} sha256 mismatch") |
| return None |
| print(f" [OK] {filename} already present") |
| return dest_path |
| if status_only: |
| print(f" [MISSING] {filename}") |
| return None |
|
|
| try: |
| from tqdm import tqdm |
| except Exception: |
| tqdm = None |
|
|
| url = f"https://zenodo.org/records/{record_id}/files/{filename}" |
| print(f" [DOWNLOAD] {filename} from Zenodo record {record_id}") |
|
|
| class ProgressHook: |
| def __init__(self) -> None: |
| self.pbar = None |
|
|
| def __call__(self, count: int, block_size: int, total_size: int) -> None: |
| if tqdm is None: |
| return |
| if self.pbar is None: |
| self.pbar = tqdm(total=total_size, unit="B", unit_scale=True, desc=filename) |
| self.pbar.update(count * block_size - self.pbar.n) |
|
|
| hook = ProgressHook() |
| try: |
| urllib.request.urlretrieve(url, dest_path, reporthook=hook) |
| if hook.pbar: |
| hook.pbar.close() |
| if expected_sha256: |
| actual = sha256_file(dest_path) |
| if actual != expected_sha256: |
| dest_path.unlink(missing_ok=True) |
| print(f" [FAIL] {filename}: sha256 mismatch") |
| return None |
| print(f" [OK] {filename} downloaded") |
| return dest_path |
| except Exception as exc: |
| if hook.pbar: |
| hook.pbar.close() |
| dest_path.unlink(missing_ok=True) |
| print(f" [FAIL] {filename}: {exc}") |
| return None |
|
|
|
|
| def prefetch_timm_weights(model_name: str, status_only: bool = False) -> bool: |
| ok, detail = import_status("timm") |
| if not ok: |
| print(f" [SKIP] timm:{model_name} - timm missing ({detail})") |
| return False |
| if status_only: |
| print(f" [CHECK] timm:{model_name} cache status cannot be proven without instantiation") |
| return True |
|
|
| import timm |
|
|
| candidates = TIMM_ALIASES.get(model_name, (model_name,)) |
| for candidate in candidates: |
| try: |
| model = timm.create_model(candidate, pretrained=True, num_classes=0) |
| del model |
| print(f" [OK] timm {model_name} via {candidate}") |
| return True |
| except RuntimeError as exc: |
| if "Unknown model" not in str(exc): |
| print(f" [FAIL] timm {model_name}: {exc}") |
| return False |
| print(f" [WARN] timm {getattr(timm, '__version__', 'unknown')} does not provide {model_name}") |
| return False |
|
|
|
|
| def prefetch_torchvision_weights(model_name: str, weight_name: str, status_only: bool = False) -> bool: |
| ok, detail = import_status("torchvision.models") |
| if not ok: |
| print(f" [SKIP] torchvision:{model_name} - torchvision missing ({detail})") |
| return False |
| if status_only: |
| print(f" [CHECK] torchvision:{model_name} cache status cannot be proven without instantiation") |
| return True |
|
|
| import torchvision.models as models |
|
|
| builder = getattr(models, model_name) |
| enum_name = "".join(part.capitalize() for part in model_name.split("_")) + "_Weights" |
| weights_enum = getattr(models, enum_name, None) |
| try: |
| if weights_enum is not None: |
| weights = getattr(weights_enum, weight_name, weights_enum.DEFAULT) |
| model = builder(weights=weights) |
| else: |
| model = builder(pretrained=True) |
| del model |
| print(f" [OK] torchvision {model_name}") |
| return True |
| except Exception as exc: |
| print(f" [FAIL] torchvision {model_name}: {exc}") |
| return False |
|
|
|
|
| def download_weights(status_only: bool = False) -> dict[str, int]: |
| print("\nPRETRAINED WEIGHTS") |
| print("------------------") |
| ok = missing = 0 |
| for spec in ZENODO_WEIGHTS.values(): |
| path = download_zenodo_weight( |
| spec["record_id"], |
| spec["filename"], |
| spec["dest_dir"], |
| expected_sha256=spec["sha256"], |
| status_only=status_only, |
| ) |
| ok += int(path is not None) |
| missing += int(path is None) |
| for model_name in TIMM_WEIGHTS_TO_PREFETCH: |
| ok += int(prefetch_timm_weights(model_name, status_only=status_only)) |
| for model_name, weight_name in TORCHVISION_WEIGHTS_TO_PREFETCH: |
| ok += int(prefetch_torchvision_weights(model_name, weight_name, status_only=status_only)) |
| return {"ok": ok, "missing": missing, "total": len(ZENODO_WEIGHTS) + len(TIMM_WEIGHTS_TO_PREFETCH) + len(TORCHVISION_WEIGHTS_TO_PREFETCH)} |
|
|
|
|
| def parse_simple_yaml(path: Path) -> dict[str, object]: |
| data: dict[str, object] = {} |
| stack: list[tuple[int, dict[str, object]]] = [(-1, data)] |
| for raw in path.read_text(encoding="utf-8").splitlines(): |
| line = raw.split("#", 1)[0].rstrip() |
| if not line.strip() or ":" not in line: |
| continue |
| indent = len(line) - len(line.lstrip(" ")) |
| key, value = line.strip().split(":", 1) |
| while stack and indent <= stack[-1][0]: |
| stack.pop() |
| current = stack[-1][1] |
| value = value.strip() |
| if not value: |
| child: dict[str, object] = {} |
| current[key] = child |
| stack.append((indent, child)) |
| continue |
| if value.startswith("[") and value.endswith("]"): |
| current[key] = [x.strip().strip("'\"") for x in value[1:-1].split(",") if x.strip()] |
| elif value.isdigit(): |
| current[key] = int(value) |
| else: |
| current[key] = value.strip("'\"") |
| return data |
|
|
|
|
| def dataset_root_from_config(cfg: dict[str, object]) -> Path: |
| value = str(cfg.get("data_root", "")) |
| data_root = os.environ.get("DATA_ROOT", str(ROOT.parent / "Datasets")) |
| return Path(value.replace("${DATA_ROOT}", data_root)).expanduser() |
|
|
|
|
| def prepare_dataset_lists(data_root: Path, image_a_folder: str = "A", splits: tuple[str, ...] = ("train", "val", "test")) -> dict[str, int]: |
| list_dir = data_root / "list" |
| list_dir.mkdir(parents=True, exist_ok=True) |
| counts: dict[str, int] = {} |
| for split in splits: |
| out = list_dir / f"{split}.txt" |
| if out.exists(): |
| counts[split] = sum(1 for _ in out.open("r", encoding="utf-8")) |
| print(f" [OK] {out} already exists ({counts[split]} entries)") |
| continue |
| split_dir = data_root / split |
| img_dir = None |
| for candidate in (image_a_folder, "A", "T1", "t1", "img", "images", "image"): |
| if (split_dir / candidate).is_dir(): |
| img_dir = split_dir / candidate |
| break |
| if img_dir is None: |
| counts[split] = 0 |
| print(f" [WARN] no image folder found under {split_dir}") |
| continue |
| names = sorted(p.name for p in img_dir.iterdir() if p.is_file() and p.suffix.lower() in IMG_EXTS) |
| out.write_text("\n".join(names) + ("\n" if names else ""), encoding="utf-8") |
| counts[split] = len(names) |
| print(f" [OK] {out}: {len(names)} entries") |
| return counts |
|
|
|
|
| def prepare_datasets(selected: str | None = None, status_only: bool = False) -> dict[str, int]: |
| print("\nDATASETS") |
| print("--------") |
| exists = missing = prepared = 0 |
| paths = sorted(DATASET_CONFIGS.glob("*.yaml")) |
| for path in paths: |
| if selected and path.stem != selected: |
| continue |
| cfg = parse_simple_yaml(path) |
| root = dataset_root_from_config(cfg) |
| if root.is_dir(): |
| exists += 1 |
| print(f" [EXISTS] {path.stem}: {root}") |
| if not status_only: |
| split_keys = tuple((cfg.get("splits") or {"train": "train", "val": "val", "test": "test"}).keys()) |
| prepare_dataset_lists(root, image_a_folder=str(cfg.get("image_a_folder", "A")), splits=split_keys) |
| prepared += 1 |
| else: |
| missing += 1 |
| print(f" [MISSING] {path.stem}: {root}") |
| return {"exists": exists, "missing": missing, "prepared": prepared, "total": len(paths)} |
|
|
|
|
| def print_final_status(env: dict[str, int], repos: dict[str, int], weights: dict[str, int] | None, datasets: dict[str, int]) -> None: |
| print("\nSETUP STATUS") |
| print("Component Status Notes") |
| print("-----------------------------|-------------|----------------------------------") |
| env_status = "COMPLETE" if env["missing"] == 0 else "CHECK" |
| print(f"Python environment {env_status:<12} {env['installed']} packages auto-installed; {env['warnings']} manual warnings") |
| print(f"Model repos ({repos['total']} total) {repos['ok']} / {repos['total']:<6} {repos['skipped']} URL-unconfirmed entries") |
| if weights is None: |
| print("Backbone weights SKIPPED run: python setup.py --weights-only") |
| else: |
| print(f"Backbone weights {weights['ok']} / {weights['total']:<6} Zenodo, timm, and torchvision warmups") |
| print(f"Datasets {datasets['exists']} / {datasets['total']:<6} roots found under DATA_ROOT/default search") |
| print("\nNext steps:") |
| print(" export DATA_ROOT=/path/to/datasets") |
| print(" python run_training.py --model bifa --dataset levir_cd --dry-run") |
| print(" python run_training.py --model bifa --dataset levir_cd") |
| print(" For ChangeMamba/Changer, install the MMSeg stack shown above.") |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description="One-command setup for the CD-Models benchmark suite.") |
| parser.add_argument("--skip-weights", action="store_true") |
| parser.add_argument("--weights-only", action="store_true") |
| parser.add_argument("--env-check-only", action="store_true") |
| parser.add_argument("--status", action="store_true") |
| parser.add_argument("--dataset", default=None, help="Prepare list files for one dataset config name.") |
| return parser.parse_args() |
|
|
|
|
| def main() -> int: |
| args = parse_args() |
| status_only = args.status |
|
|
| env = check_environment(status_only=status_only) |
| if args.env_check_only: |
| print_final_status(env, {"ok": 0, "missing": 0, "skipped": 0, "total": len(CONFIRMED_REPOS)}, None, {"exists": 0, "missing": 0, "prepared": 0, "total": 0}) |
| return 0 if env["warnings"] == 0 else 1 |
|
|
| repos = {"ok": 0, "missing": 0, "skipped": 0, "total": len(CONFIRMED_REPOS)} |
| if not args.weights_only: |
| repos = clone_repositories(status_only=status_only) |
|
|
| weights = None |
| if not args.skip_weights: |
| weights = download_weights(status_only=status_only) |
|
|
| datasets = prepare_datasets(selected=args.dataset, status_only=status_only) |
| print_final_status(env, repos, weights, datasets) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
