from __future__ import annotations import argparse import json import os import shutil import subprocess import sys from pathlib import Path import torch ROOT = Path(__file__).resolve().parents[1] if str(ROOT) not in sys.path: sys.path.insert(0, str(ROOT)) from utils.gpu_utils import ( gpu_mapping_message, print_gpu_diagnostics, resolve_gpu, subprocess_gpu_env, ) from utils.dataset_cache import apply_dataloader_cli_overrides, dataset_runtime_summary, print_dataloader_policy SKIP_EXIT_CODE = 75 FAILED_SUBPROCESS_EXIT_CODE = 70 FAILED_NO_CHECKPOINT_EXIT_CODE = 71 FAILED_EVAL_EXIT_CODE = 72 EXISTING_LEGACY = { "bifa": ("BiFA", ["python3", "train_wildfire.py"]), "bit_cd": ("BIT_CD", ["python3", "train_wildfire.py"]), "cdmamba": ("CDMamba", ["python3", "train_wildfire.py"]), "change3d": ("Change3D", ["python3", "train_wildfire.py"]), "changeformer": ("ChangeFormer", ["python3", "train_wildfire.py"]), "dsifn": ("IFNet", ["python3", "train_wildfire.py"]), "ifnet": ("IFNet", ["python3", "train_wildfire.py"]), "rsm_cd": ("RSM-CD/change_detection_mamba", ["python3", "train_wildfire.py"]), "schanger": ("SChanger", ["python3", "train_wildfire.py"]), "siam_nestedunet": ("Siam-NestedUNet", ["python3", "train_wildfire.py"]), "stanet": ("STANet", ["python3", "train_wildfire.py"]), } JSON_LEGACY_TRAINERS = {"bifa", "cdmamba"} MAIN_CD_TRAINERS = {"bit_cd", "changeformer"} NEW_REPOS = { "fc_ef": ("model_repos/fully_convolutional_change_detection", "https://github.com/rcdaudt/fully_convolutional_change_detection"), "fc_siam_conc": ("model_repos/fully_convolutional_change_detection", "https://github.com/rcdaudt/fully_convolutional_change_detection"), "fc_siam_diff": ("model_repos/fully_convolutional_change_detection", "https://github.com/rcdaudt/fully_convolutional_change_detection"), "dsifn": ("model_repos/DSIFN", "original DSIFN repo still requires URL confirmation"), "changemamba": ("model_repos/ChangeMamba", "https://github.com/ChenHongruixuan/ChangeMamba"), "elgcnet": ("model_repos/elgcnet", "https://github.com/techmn/elgcnet"), "changer": ("model_repos/open-cd", "https://github.com/likyoo/open-cd"), "hanet": ("model_repos/HANet-CD", "https://github.com/ChengxiHAN/HANet-CD"), "cgnet": ("model_repos/CGNet-CD", "https://github.com/ChengxiHAN/CGNet-CD"), "dsamnet": ("model_repos/DSAMNet", "https://github.com/liumency/DSAMNet"), "tinycd": ("model_repos/Tiny_model_4_CD", "https://github.com/AndreaCodegoni/Tiny_model_4_CD"), } TORCHVISION_WEIGHTS = { "changer": "resnet18", "dsamnet": "resnet18", "hanet": "resnet50", "cgnet": "resnet50", "dsifn": "vgg16", } TIMM_WEIGHTS = { "tinycd": "efficientnet_b4", "elgcnet": "mit_b0", "changeformer": "mit_b1", } EXTERNAL_TRAINABLE = {"dsamnet", "cgnet", "hanet", "tinycd"} def parse_args(model_name: str) -> argparse.Namespace: parser = argparse.ArgumentParser(description=f"Train/evaluate {model_name} on a cd-models dataset.") parser.add_argument("--dataset", default="wildfire_s2") parser.add_argument("--epochs", type=int, default=None) parser.add_argument("--batch-size", type=int, default=None) parser.add_argument("--lr", type=float, default=None) parser.add_argument("--gpu", default="0") parser.add_argument("--resume", action="store_true") parser.add_argument("--eval-only", action="store_true") parser.add_argument("--dry-run", action="store_true") parser.add_argument("--smoke-test", action="store_true") parser.add_argument("--force", action="store_true") parser.add_argument("--output-dir", default=None) parser.add_argument("--model", default=model_name, help="Variant selector used by FC-Siam wrappers.") parser.add_argument("--max-iters", type=int, default=None, help="Iteration override for iteration-based upstream trainers.") parser.add_argument("--num-workers", type=int, default=None) parser.add_argument("--prefetch-factor", type=int, default=None) parser.set_defaults(persistent_workers=None, pin_memory=None) parser.add_argument("--persistent-workers", dest="persistent_workers", action="store_true") parser.add_argument("--no-persistent-workers", dest="persistent_workers", action="store_false") parser.add_argument("--pin-memory", dest="pin_memory", action="store_true") parser.add_argument("--no-pin-memory", dest="pin_memory", action="store_false") return parser.parse_args() def _ensure_pretrained(model_name: str, variant: str = "tiny") -> str | None: from utils.weight_downloader import ensure_timm_weight, ensure_torchvision_weight, ensure_weights if model_name == "changemamba": return ensure_weights(f"vmamba_{variant}") if model_name in TIMM_WEIGHTS: ensure_timm_weight(TIMM_WEIGHTS[model_name], required=model_name == "tinycd") if model_name in TORCHVISION_WEIGHTS: ensure_torchvision_weight(TORCHVISION_WEIGHTS[model_name]) return None def _dataset_ready(dataset_name: str, return_format: str = "tuple", dry_run: bool = False) -> dict: from datasets.cd_dataset import CDDataset from utils.config_loader import load_dataset_config cfg = load_dataset_config(dataset_name) root_exists = Path(cfg["data_root"]).exists() print( f"[DATASET] {cfg['name']} root={cfg['data_root']} img_size={cfg.get('img_size')} " f"channels={cfg.get('channels')} return_format={return_format}" ) print(f"[DATASET] {dataset_runtime_summary(cfg)}") print_dataloader_policy(cfg, torch.cuda.is_available()) if cfg.get("io_warning"): print(f"[DATASET-WARNING] {cfg['io_warning']}") print(f"[DRY-RUN] Dataset root exists: {root_exists} -> {cfg['data_root']}" if dry_run else f"[DATASET] Root exists: {root_exists}") if dry_run: return cfg if not root_exists: raise FileNotFoundError( f"Dataset root directory not found for {dataset_name}: {cfg['data_root']}. " "Set DATA_ROOT or edit the dataset YAML." ) try: from utils.dataset_list_generator import generate_list_files list_dir = Path(cfg["data_root"]) / "list" if not (list_dir / "train.txt").exists(): print(f"[SETUP] Generating list files for {dataset_name}...") generate_list_files( cfg["data_root"], splits=list((cfg.get("splits") or {"train": "train", "val": "val", "test": "test"}).keys()), img_subdir=cfg.get("image_a_folder", "A"), ) except Exception as exc: print(f"[SETUP] List-file generation skipped: {exc}") if not Path(cfg["data_root"]).exists(): print(f"[DATASET] Root does not exist yet, skipping CDDataset construction: {cfg['data_root']}") return cfg _ = CDDataset(cfg["data_root"], "train", cfg=cfg, return_format=return_format) return cfg def _run_logged(cmd: list[str], cwd: Path, env: dict, stdout_path: Path, stderr_path: Path) -> int: stdout_path.parent.mkdir(parents=True, exist_ok=True) stderr_path.parent.mkdir(parents=True, exist_ok=True) visible = env.get("CUDA_VISIBLE_DEVICES") if visible: print(_gpu_mapping_message(env)) with stdout_path.open("w", encoding="utf-8") as stdout, stderr_path.open("w", encoding="utf-8") as stderr: return subprocess.run(cmd, cwd=cwd, env=env, stdout=stdout, stderr=stderr, check=False).returncode def _run_streamed(cmd: list[str], cwd: Path, env: dict, log_path: Path) -> int: log_path.parent.mkdir(parents=True, exist_ok=True) with log_path.open("a", encoding="utf-8") as log: log.write("\n[WRAPPER] cwd=" + str(cwd) + "\n") log.write("[WRAPPER] " + _gpu_mapping_message(env) + "\n") log.write("[WRAPPER] command=" + " ".join(cmd) + "\n") log.flush() process = subprocess.Popen( cmd, cwd=cwd, env=env, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, bufsize=1, ) assert process.stdout is not None for line in process.stdout: print(line, end="", flush=True) log.write(line) log.flush() return process.wait() def _gpu_mapping_message(env: dict) -> str: return gpu_mapping_message(env=env) def _print_train_command(cwd: Path, cmd: list[str], env: dict) -> None: print(_gpu_mapping_message(env)) print("[TRAIN]", cwd, " ".join(cmd)) def _evaluate(model_name: str, dataset: str, gpu: str, dry_run: bool) -> int: cmd = [sys.executable, str(ROOT / "evaluate.py"), "--model", model_name, "--dataset", dataset, "--gpu", gpu] print("[EVAL]", " ".join(cmd)) if dry_run: return 0 log_dir = ROOT / "results" / model_name / dataset / "logs" env, _ = subprocess_gpu_env(gpu) return _run_logged(cmd, ROOT, env, log_dir / "eval_stdout.log", log_dir / "eval_stderr.log") def _legacy_checkpoint_candidates(model_name: str, dataset_name: str) -> list[Path]: if model_name != "cdmamba": return [] experiments = ROOT / "CDMamba" / "experiments" if not experiments.exists(): return [] candidates: list[Path] = [] run_dirs = sorted( ( path for path in experiments.glob(f"{dataset_name}-train-cdmamba_*") if path.is_dir() ), key=lambda path: path.stat().st_mtime, ) preferred_patterns = [ "checkpoint/best_cd_model_gen.pth", "checkpoint/*_gen.pth", "checkpoint/*best*.pth", "checkpoint/*.pth", "**/*.pth", "**/*.pt", ] for run_dir in run_dirs: for pattern in preferred_patterns: matches = [ path for path in sorted(run_dir.glob(pattern), key=lambda item: item.stat().st_mtime) if path.is_file() and "_opt" not in path.name.lower() ] if matches: candidates.extend(matches) break return candidates def _canonicalize_checkpoints(model_name: str, dataset_name: str) -> None: out_dir = ROOT / "results" / model_name / dataset_name ckpt_dir = out_dir / "checkpoints" ckpt_dir.mkdir(parents=True, exist_ok=True) best_out = ckpt_dir / "best_model.pth" latest_out = ckpt_dir / "latest.pth" if best_out.exists(): return patterns = [ "**/best_model.pth", "**/best_ckpt.pth", "**/best_ckpt.pt", "**/*best*.pth", "**/*best*.pt", "**/*_best_iou.pth", "**/netCD_epoch_*.pth", "**/*.pth", "**/*.pt", ] candidates: list[Path] = [] for pattern in patterns: for path in sorted(out_dir.glob(pattern)): if path.is_file() and path.resolve() != best_out.resolve(): candidates.append(path) if candidates: break if not candidates: candidates = _legacy_checkpoint_candidates(model_name, dataset_name) if not candidates: raise FileNotFoundError( f"Training completed for {model_name}/{dataset_name}, but no checkpoint file was found under " f"{out_dir} or known legacy output directories." ) source = candidates[-1] shutil.copy2(source, best_out) if not latest_out.exists(): shutil.copy2(source, latest_out) metadata = { "model": model_name, "dataset": dataset_name, "canonical_best_model": str(best_out), "source_checkpoint": str(source), "source_checkpoint_name": source.name, } (ckpt_dir / "checkpoint_metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") def _checkpoint_files(checkpoint_dir: Path) -> list[Path]: suffixes = {".pth", ".pt", ".ckpt"} if not checkpoint_dir.exists(): return [] return sorted(path for path in checkpoint_dir.rglob("*") if path.is_file() and path.suffix.lower() in suffixes) def _print_checkpoint_failure(model_name: str, dataset_name: str, cmd: list[str], checkpoint_dir: Path, result_dir: Path) -> None: print(f"[FAILED_NO_CHECKPOINT] {model_name}/{dataset_name}: upstream command returned successfully but no checkpoint was found.") print(f"[FAILED_NO_CHECKPOINT] command={' '.join(cmd)}") print(f"[FAILED_NO_CHECKPOINT] expected checkpoint dir={checkpoint_dir}") if result_dir.exists(): existing = [str(path.relative_to(result_dir)) for path in sorted(result_dir.rglob("*"))[:80]] print(f"[FAILED_NO_CHECKPOINT] existing result files={existing}") else: print(f"[FAILED_NO_CHECKPOINT] result dir does not exist: {result_dir}") def _iteration_budget(model_name: str, model_cfg: dict, args: argparse.Namespace, default: int) -> tuple[int, str]: if args.max_iters is not None: return int(args.max_iters), "cli --max-iters" if args.smoke_test: return int(model_cfg.get("smoke_max_iters", min(default, 2))), "model smoke_max_iters" if "max_iters" in model_cfg: return int(model_cfg["max_iters"]), "model max_iters" return default, f"{model_name} wrapper default" def _model_config(model_name: str) -> dict: from utils.config_loader import load_model_config return load_model_config(model_name) def _legacy_dataset_token(dataset_cfg: dict) -> str: name = str(dataset_cfg.get("name", "")).lower() source = str(dataset_cfg.get("source_name", "")).lower() if "dsifn" in name or "dsifn" in source: return "DSIFN" if "levir" in name or "levir" in source: return "LEVIR" if "whu" in name or "whu" in source: return "WHU" if "wildfire" in name or "wildfire" in source: return "WildFireS2" return dataset_cfg.get("source_name", dataset_cfg["name"]) def _split_dir(dataset_cfg: dict, split: str) -> Path: return Path(dataset_cfg["data_root"]) / dataset_cfg.get("splits", {}).get(split, split) def _safe_link_or_copy(src: Path, dst: Path) -> None: if dst.is_symlink(): try: if Path(os.readlink(dst)) == src: return except OSError: pass dst.unlink() if dst.exists(): return dst.parent.mkdir(parents=True, exist_ok=True) if not src.exists(): print(f"[DATASET-VIEW] Source path is not present yet, skipping generated link: {src}") return try: os.symlink(src, dst, target_is_directory=src.is_dir()) except OSError: if src.is_dir(): shutil.copytree(src, dst, dirs_exist_ok=True) else: shutil.copy2(src, dst) def _safe_mask_link_or_copy(src: Path, dst: Path) -> None: dst.parent.mkdir(parents=True, exist_ok=True) if not src.exists(): print(f"[DATASET-VIEW] Source mask is not present yet, skipping generated link: {src}") return try: import numpy as np from PIL import Image arr = np.asarray(Image.open(src)) if arr.ndim == 3: arr = arr[..., 0] if arr.size and int(arr.max()) <= 1: if dst.exists() or dst.is_symlink(): dst.unlink() Image.fromarray((arr > 0).astype(np.uint8) * 255).save(dst, format="PNG") return except Exception as exc: print(f"[DATASET-VIEW] Could not inspect mask {src}: {exc}") _safe_link_or_copy(src, dst) def _scan_images(folder: Path) -> dict[str, Path]: exts = {".png", ".jpg", ".jpeg", ".tif", ".tiff", ".bmp"} if not folder.is_dir(): return {} return { p.stem: p for p in sorted(folder.iterdir()) if p.is_file() and p.suffix.lower() in exts } def _view_split_counts(dataset_cfg: dict) -> dict[str, int]: counts: dict[str, int] = {} for split in ("train", "val", "test"): split_root = _split_dir(dataset_cfg, split) a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) counts[split] = len(set(a_files) & set(b_files) & set(label_files)) return counts def _log_generated_view(dataset_cfg: dict, view: Path) -> None: counts = _view_split_counts(dataset_cfg) print(f"[DATASET] dataset name: {dataset_cfg['name']}") print(f"[DATASET] original root: {dataset_cfg.get('original_data_root', dataset_cfg.get('data_root'))}") print(f"[DATASET] local root: {dataset_cfg.get('local_root')}") print(f"[DATASET] using local root: {bool(dataset_cfg.get('using_local_root'))}") print(f"[DATASET] generated view path: {view}") print(f"[DATASET] train/val/test counts: {counts}") def _prepare_matched_split_view(dataset_cfg: dict, split: str) -> Path: split_root = _split_dir(dataset_cfg, split) view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "matched" / split a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) for folder in ("A", "B", "label"): (view / folder).mkdir(parents=True, exist_ok=True) for stem in sorted(set(a_files) & set(b_files) & set(label_files)): target_name = a_files[stem].name _safe_link_or_copy(a_files[stem], view / "A" / target_name) _safe_link_or_copy(b_files[stem], view / "B" / target_name) _safe_mask_link_or_copy(label_files[stem], view / "label" / target_name) _log_generated_view(dataset_cfg, view) return view def _prepare_split_view(dataset_cfg: dict, split: str) -> Path: split_root = _split_dir(dataset_cfg, split) view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / split mapping = { "A": split_root / dataset_cfg.get("image_a_folder", "A"), "B": split_root / dataset_cfg.get("image_b_folder", "B"), "label": split_root / dataset_cfg.get("mask_folder", "label"), } for target, src in mapping.items(): _safe_link_or_copy(src, view / target) _log_generated_view(dataset_cfg, view) return view def _prepare_tinycd_view(dataset_cfg: dict) -> Path: view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "tinycd" list_dir = view / "list" list_dir.mkdir(parents=True, exist_ok=True) for folder in ("A", "B", "label"): (view / folder).mkdir(parents=True, exist_ok=True) for split in ("train", "val", "test"): split_root = _split_dir(dataset_cfg, split) a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) names = [] for stem in sorted(set(a_files) & set(b_files) & set(label_files)): target_name = f"{split}__{a_files[stem].name}" _safe_link_or_copy(a_files[stem], view / "A" / target_name) _safe_link_or_copy(b_files[stem], view / "B" / target_name) _safe_mask_link_or_copy(label_files[stem], view / "label" / target_name) names.append(target_name) (list_dir / f"{split}.txt").write_text("\n".join(names) + ("\n" if names else ""), encoding="utf-8") _log_generated_view(dataset_cfg, view) return view def _prepare_changemamba_view(dataset_cfg: dict) -> Path: view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "changemamba" split_map = {"train": "train", "test": "val"} for target_split, source_split in split_map.items(): for folder in ("T1", "T2", "GT"): (view / target_split / folder).mkdir(parents=True, exist_ok=True) split_root = _split_dir(dataset_cfg, source_split) a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) names = [] for stem in sorted(set(a_files) & set(b_files) & set(label_files)): target_name = f"{source_split}__{a_files[stem].name}" _safe_link_or_copy(a_files[stem], view / target_split / "T1" / target_name) _safe_link_or_copy(b_files[stem], view / target_split / "T2" / target_name) _safe_mask_link_or_copy(label_files[stem], view / target_split / "GT" / target_name) names.append(target_name) list_name = "train_set.txt" if target_split == "train" else "test_set.txt" (view / list_name).write_text("\n".join(names) + ("\n" if names else ""), encoding="utf-8") _log_generated_view(dataset_cfg, view) return view def _prepare_opencd_view(dataset_cfg: dict) -> Path: view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "opencd" for split in ("train", "val", "test"): for folder in ("A", "B", "label"): (view / split / folder).mkdir(parents=True, exist_ok=True) split_root = _split_dir(dataset_cfg, split) a_files = _scan_images(split_root / dataset_cfg.get("image_a_folder", "A")) b_files = _scan_images(split_root / dataset_cfg.get("image_b_folder", "B")) label_files = _scan_images(split_root / dataset_cfg.get("mask_folder", "label")) for stem in sorted(set(a_files) & set(b_files) & set(label_files)): image_name = a_files[stem].name label_name = f"{a_files[stem].stem}{label_files[stem].suffix}" _safe_link_or_copy(a_files[stem], view / split / "A" / image_name) _safe_link_or_copy(b_files[stem], view / split / "B" / image_name) _safe_link_or_copy(label_files[stem], view / split / "label" / label_name) _log_generated_view(dataset_cfg, view) return view def _write_opencd_changer_config(dataset_cfg: dict, model_cfg: dict, view: Path) -> Path: train_split = _split_dir(dataset_cfg, "train") n_train = len(_scan_images(train_split / dataset_cfg.get("image_a_folder", "A"))) batch_size = int(dataset_cfg.get("batch_size", 8)) steps_per_epoch = max((n_train + batch_size - 1) // batch_size, 1) max_iters = int(model_cfg.get("max_iters", int(model_cfg.get("num_epochs", 200)) * steps_per_epoch)) val_interval = max(steps_per_epoch, 1) cfg_path = ROOT / "generated_configs" / f"{dataset_cfg['name']}__changer_opencd.py" cfg_path.parent.mkdir(parents=True, exist_ok=True) cfg_path.write_text( "\n".join( [ f"_base_ = '{ROOT / 'model_repos' / 'open-cd' / 'configs' / 'changer' / 'changer_ex_r18_512x512_40k_levircd.py'}'", "", "dataset_type = 'DSIFN_Dataset'", f"data_root = r'{view}'", f"crop_size = ({int(dataset_cfg.get('img_size', model_cfg.get('img_size', 256)))}, {int(dataset_cfg.get('img_size', model_cfg.get('img_size', 256)))})", "", "train_dataloader = dict(", f" batch_size={batch_size},", f" num_workers={int(dataset_cfg.get('num_workers', 4))},", " dataset=dict(", " type=dataset_type,", " data_root=data_root,", " data_prefix=dict(seg_map_path='train/label', img_path_from='train/A', img_path_to='train/B')))", "val_dataloader = dict(", " dataset=dict(", " type=dataset_type,", " data_root=data_root,", " data_prefix=dict(seg_map_path='val/label', img_path_from='val/A', img_path_to='val/B')))", "test_dataloader = dict(", " dataset=dict(", " type=dataset_type,", " data_root=data_root,", " data_prefix=dict(seg_map_path='test/label', img_path_from='test/A', img_path_to='test/B')))", "", f"train_cfg = dict(type='IterBasedTrainLoop', max_iters={max_iters}, val_interval={val_interval})", f"work_dir = r'{ROOT / 'results' / 'changer' / dataset_cfg['name'] / 'work_dir'}'", "", ] ), encoding="utf-8", ) return cfg_path def _write_hanet_metadata(dataset_cfg: dict, model_cfg: dict, view_root: Path) -> Path: out_dir = ROOT / "results" / "hanet" / dataset_cfg["name"] metadata = { "patch_size": int(dataset_cfg.get("img_size", 256)), "augmentation": True, "num_gpus": 1, "num_workers": int(dataset_cfg.get("num_workers", 4)), "num_channel": 3, "EF": False, "epochs": int(model_cfg.get("num_epochs", 50)), "epochs_threshold": 15, "gamma": 0.5, "weight_decay": float(model_cfg.get("weight_decay", 5e-4) or 5e-4), "batch_size": int(dataset_cfg.get("batch_size", 8)), "learning_rate": float(model_cfg.get("lr", 5e-4)), "loss_function": "hybrid", "dataset_dir": str(view_root) + "/", "weight_dir": str(out_dir / "weights") + "/", "Output_dir": str(out_dir / "outputs") + "/", "log_dir": str(out_dir / "logs"), } path = ROOT / "generated_configs" / f"{dataset_cfg['name']}__hanet_metadata.json" path.parent.mkdir(parents=True, exist_ok=True) path.write_text(json.dumps(metadata, indent=2), encoding="utf-8") return path def _existing_train_command(model_name: str, dataset_cfg: dict, args: argparse.Namespace) -> tuple[Path, list[str], dict]: workdir_rel, default_cmd = EXISTING_LEGACY[model_name] workdir = ROOT / workdir_rel env, _ = subprocess_gpu_env(args.gpu) child_gpu = "0" env["CD_MODELS_DATASET_ROOT"] = dataset_cfg["data_root"] env["CD_MODELS_DSIFN_ROOT"] = dataset_cfg["data_root"] env["CD_MODELS_DATASET_NAME"] = dataset_cfg["name"] env["CD_MODELS_NUM_WORKERS"] = str(int(dataset_cfg.get("num_workers", 2))) env["CD_MODELS_PREFETCH_FACTOR"] = str(int(dataset_cfg.get("prefetch_factor", 2))) env["CD_MODELS_PERSISTENT_WORKERS"] = "1" if dataset_cfg.get("persistent_workers") else "0" env["CD_MODELS_PIN_MEMORY"] = "1" if dataset_cfg.get("pin_memory", True) else "0" env["CD_MODELS_BATCH_SIZE"] = str(int(args.batch_size or dataset_cfg.get("batch_size", 8))) if args.epochs is not None: env["CD_MODELS_EPOCHS"] = str(int(args.epochs)) if dataset_cfg["name"] == "wildfire_s2": return workdir, [sys.executable, default_cmd[1]], env if model_name in JSON_LEGACY_TRAINERS: from utils.legacy_config_writer import write_bifa_or_cdmamba_config model_cfg = _model_config(model_name) if args.epochs is not None: model_cfg["num_epochs"] = int(args.epochs) if args.lr is not None: model_cfg["lr"] = float(args.lr) if args.batch_size is not None: dataset_cfg["batch_size"] = int(args.batch_size) config_path = write_bifa_or_cdmamba_config( model_name, dataset_cfg, model_cfg, prepare_view=not args.dry_run, ) return workdir, [sys.executable, "train_cd.py", "--config", str(config_path), "--phase", "train", "--gpu_ids", child_gpu], env if model_name in MAIN_CD_TRAINERS: from utils.legacy_config_writer import prepare_legacy_list_view legacy_root = ( ROOT / "generated_dataset_views" / dataset_cfg["name"] / "legacy_list" if args.dry_run else prepare_legacy_list_view(dataset_cfg) ) env["CD_MODELS_DATASET_ROOT"] = str(legacy_root) env["CD_MODELS_DSIFN_ROOT"] = str(legacy_root) model_cfg = _model_config(model_name) token = _legacy_dataset_token(dataset_cfg) project = f"{dataset_cfg['name']}-train-{model_name}" if model_name == "bit_cd": net_g = "base_transformer_pos_s4_dd8" optimizer = "sgd" loss = "ce" lr = float(model_cfg.get("lr", 0.01)) else: net_g = "ChangeFormerV6" optimizer = "adamw" loss = "ce" lr = float(model_cfg.get("lr", 0.00006)) cmd = [ sys.executable, "main_cd.py", "--gpu_ids", child_gpu, "--project_name", project, "--data_name", token, "--img_size", str(int(dataset_cfg.get("img_size", model_cfg.get("img_size", 256)))), "--batch_size", str(int(dataset_cfg.get("batch_size", 8))), "--num_workers", str(int(dataset_cfg.get("num_workers", 4))), "--max_epochs", str(int(model_cfg.get("num_epochs", 200))), "--optimizer", optimizer, "--lr", str(lr), "--loss", loss, "--net_G", net_g, ] return workdir, cmd, env from utils.legacy_config_writer import prepare_legacy_list_view legacy_root = ( ROOT / "generated_dataset_views" / dataset_cfg["name"] / "legacy_list" if args.dry_run else prepare_legacy_list_view(dataset_cfg) ) out_dir = ROOT / "results" / model_name / dataset_cfg["name"] env["CD_MODELS_DATASET_ROOT"] = str(legacy_root) env["CD_MODELS_DSIFN_ROOT"] = str(legacy_root) env["CD_MODELS_DATASET_NAME"] = dataset_cfg["name"] env["CD_MODELS_CHECKPOINT_DIR"] = str(out_dir / "checkpoints") env["CD_MODELS_LOG_PATH"] = str(out_dir / "logs" / "train.log") if args.lr is not None: env["CD_MODELS_LR"] = str(float(args.lr)) else: env["CD_MODELS_LR"] = str(float(_model_config(model_name).get("lr", 1e-3))) return workdir, [sys.executable, "train_wildfire.py"], env def run_existing_model(model_name: str, args: argparse.Namespace) -> int: from utils.model_adapters import get_model_adapter from utils.unified_trainer import train_with_adapter adapter = get_model_adapter(model_name) if adapter.supports_unified_training: return train_with_adapter(model_name, args) dataset_cfg = _dataset_ready(args.dataset, return_format="legacy", dry_run=args.dry_run) apply_dataloader_cli_overrides(dataset_cfg, args) print_dataloader_policy(dataset_cfg, torch.cuda.is_available()) workdir_rel, cmd = EXISTING_LEGACY[model_name] workdir, cmd, env = _existing_train_command(model_name, dataset_cfg, args) if not cmd: print( f"[SKIP] {model_name}/{args.dataset}: no verified non-WildFire training command is wired for " f"{workdir_rel} yet." ) return 0 if args.dry_run else SKIP_EXIT_CODE _print_train_command(workdir, cmd, env) if args.dry_run: return _evaluate(model_name, args.dataset, args.gpu, dry_run=True) code = 0 if not args.eval_only: log_dir = ROOT / "results" / model_name / dataset_cfg["name"] / "logs" code = _run_logged(cmd, workdir, env, log_dir / "train_stdout.log", log_dir / "train_stderr.log") if code == 0: _canonicalize_checkpoints(model_name, dataset_cfg["name"]) code = _evaluate(model_name, args.dataset, args.gpu, dry_run=False) return code def run_external_model(model_name: str, args: argparse.Namespace) -> int: from utils.model_adapters import get_model_adapter from utils.unified_trainer import train_with_adapter adapter = get_model_adapter(model_name) if adapter.supports_unified_training: return train_with_adapter(model_name, args) cfg = _dataset_ready(args.dataset, return_format="tuple", dry_run=args.dry_run) apply_dataloader_cli_overrides(cfg, args) print_dataloader_policy(cfg, torch.cuda.is_available()) repo_rel, repo_url = NEW_REPOS[model_name] repo = ROOT / repo_rel if not repo.exists(): print( f"[SKIP] {model_name}/{args.dataset}: source repo is not present at {repo}. " f"Expected source: {repo_url}." ) return 0 if args.dry_run else SKIP_EXIT_CODE if model_name == "tinycd" and int(cfg.get("channels", 3)) > 3: raise ValueError("TinyCD supports RGB input only. Set channels: 3 in the dataset config override.") weight_path = None if not args.dry_run: weight_path = _ensure_pretrained(model_name) print(f"[READY] {model_name} repo found at {repo}") if weight_path: print(f"[WEIGHTS] {weight_path}") train_cmd, train_env, train_cwd = _external_train_command(model_name, cfg, args) if not train_cmd: print("[NOTE] Model-specific training command construction must be completed from the cloned repo config.") print(f"[SKIP] {model_name}/{args.dataset}: repo-specific train command wiring is not implemented yet.") return 0 if args.dry_run else SKIP_EXIT_CODE if args.dry_run or args.eval_only: _print_train_command(train_cwd, train_cmd, train_env) return _evaluate(model_name, args.dataset, args.gpu, dry_run=args.dry_run) _print_train_command(train_cwd, train_cmd, train_env) log_dir = ROOT / "results" / model_name / cfg["name"] / "logs" if model_name == "changemamba": result_dir = ROOT / "results" / model_name / cfg["name"] checkpoint_dir = result_dir / "checkpoints" code = _run_streamed(train_cmd, train_cwd, train_env, log_dir / "train_wrapper.log") if code != 0: return FAILED_SUBPROCESS_EXIT_CODE checkpoints = _checkpoint_files(checkpoint_dir) if not checkpoints: _print_checkpoint_failure(model_name, cfg["name"], train_cmd, checkpoint_dir, result_dir) return FAILED_NO_CHECKPOINT_EXIT_CODE _canonicalize_checkpoints(model_name, cfg["name"]) code = _evaluate(model_name, args.dataset, args.gpu, dry_run=False) return 0 if code == 0 else FAILED_EVAL_EXIT_CODE code = _run_logged(train_cmd, train_cwd, train_env, log_dir / "train_stdout.log", log_dir / "train_stderr.log") if code == 0 and model_name not in {"fc_ef", "fc_siam_conc", "fc_siam_diff"}: _canonicalize_checkpoints(model_name, cfg["name"]) code = _evaluate(model_name, args.dataset, args.gpu, dry_run=False) return code def run_smoke_test(model_name: str, args: argparse.Namespace) -> int: from datasets.cd_dataset import CDDataset from utils.config_loader import load_dataset_config, load_model_config from utils.metrics import BinaryMetrics from utils.model_adapters import get_model_adapter from utils.profiling import count_flops, count_parameters from utils.qualitative import safe_sample_id, save_binary_prediction, save_probability_map, select_or_load_manifest cfg = load_dataset_config(args.dataset) apply_dataloader_cli_overrides(cfg, args) print_dataloader_policy(cfg, torch.cuda.is_available()) root = Path(cfg["data_root"]) print(f"[SMOKE] {model_name}/{args.dataset}") print(f"[SMOKE] Dataset root exists: {root.is_dir()} -> {root}") if not root.is_dir(): return 1 ds = CDDataset(root, "test", cfg=cfg, return_format="tuple") a, b, mask, name = ds[0] print(f"[SMOKE] Batch sample: name={name} A={tuple(a.shape)} B={tuple(b.shape)} mask={tuple(mask.shape)}") adapter = get_model_adapter(model_name) if adapter.supports_inprocess_eval: gpu_resolution = resolve_gpu(args.gpu) print_gpu_diagnostics(gpu_resolution) device = torch.device(gpu_resolution.local_device) model = adapter.build_model(load_model_config(model_name), cfg, device) ckpt = ROOT / "results" / model_name / cfg["name"] / "checkpoints" / "best_model.pth" if ckpt.exists(): adapter.load_checkpoint(model, ckpt, device) print(f"[SMOKE] Loaded checkpoint: {ckpt}") else: print(f"[SMOKE] No checkpoint found at {ckpt}; validating construction and forward pass only.") batch = (a.unsqueeze(0), b.unsqueeze(0), mask.unsqueeze(0), [name]) train_step = None if adapter.supports_unified_training: model.train() optimizer = adapter.build_optimizer(model, load_model_config(model_name)) optimizer.zero_grad(set_to_none=True) raw_train = adapter.forward(model, batch, device) loss_dict = adapter.compute_loss(raw_train, batch, load_model_config(model_name), cfg, device) loss_dict["loss"].backward() optimizer.step() train_step = {key: float(value.detach().cpu().item()) for key, value in loss_dict.items()} model.eval() with torch.inference_mode(): raw = adapter.forward(model, batch, device) normalized = adapter.normalize_output(raw, batch, cfg) metrics = BinaryMetrics(threshold=float(cfg.get("eval", {}).get("threshold", 0.5))) metrics.update(normalized.metric_tensor, mask.unsqueeze(0)) params = count_parameters(model) flops = None if adapter.supports_flops: flops = count_flops(model, lambda: adapter.get_dummy_inputs(cfg, device), device) manifest = select_or_load_manifest(cfg) smoke_dir = ROOT / "results" / model_name / cfg["name"] / "smoke_test" clean_id = safe_sample_id(str(name)) save_binary_prediction(normalized.binary[0], smoke_dir / f"{clean_id}_pred.png") if normalized.score is not None: save_probability_map(normalized.score[0], smoke_dir / f"{clean_id}_prob.png") print( f"[SMOKE] PASS {model_name}/{cfg['name']} adapter={adapter.model_class_path} " f"output={tuple(normalized.metric_tensor.shape)} params_m={params['params_m']:.3f} " f"flops_g={None if flops is None else flops.get('flops_g')} " f"manifest_samples={manifest.get('selected_count')} train_step={train_step} metrics={metrics.compute()}" ) return 0 args.dry_run = True if model_name in EXISTING_LEGACY: workdir, cmd, env = _existing_train_command(model_name, cfg, args) else: cmd, env, workdir = _external_train_command(model_name, cfg, args) if cmd: print(_gpu_mapping_message(env)) print(f"[SMOKE] Command construction OK: cwd={workdir} cmd={' '.join(cmd)}") else: print("[SMOKE] Command construction reports this adapter is not wired yet.") print( f"[SMOKE-UNAVAILABLE] {model_name}/{args.dataset}: {adapter.notes_or_failure_reason}" ) return 1 def _external_train_command(model_name: str, dataset_cfg: dict, args: argparse.Namespace) -> tuple[list[str], dict, Path]: from utils.legacy_config_writer import prepare_legacy_list_view env, _ = subprocess_gpu_env(args.gpu) child_gpu = "0" model_cfg = _model_config(model_name) out_dir = ROOT / "results" / model_name / dataset_cfg["name"] out_dir.mkdir(parents=True, exist_ok=True) if model_name == "dsamnet": repo = ROOT / "model_repos" / "DSAMNet" train_root = _split_dir(dataset_cfg, "train") val_root = _split_dir(dataset_cfg, "val") ckpt_dir = out_dir / "checkpoints" ckpt_dir.mkdir(parents=True, exist_ok=True) cmd = [ sys.executable, "train.py", "--num_epochs", str(int(model_cfg.get("num_epochs", 100))), "--batchsize", str(int(dataset_cfg.get("batch_size", 8))), "--val_batchsize", str(int(dataset_cfg.get("batch_size", 8))), "--num_workers", str(int(dataset_cfg.get("num_workers", 4))), "--gpu_id", child_gpu, "--train1_dir", str(train_root / dataset_cfg.get("image_a_folder", "A")), "--train2_dir", str(train_root / dataset_cfg.get("image_b_folder", "B")), "--label_train", str(train_root / dataset_cfg.get("mask_folder", "label")), "--val1_dir", str(val_root / dataset_cfg.get("image_a_folder", "A")), "--val2_dir", str(val_root / dataset_cfg.get("image_b_folder", "B")), "--label_val", str(val_root / dataset_cfg.get("mask_folder", "label")), "--model_dir", str(ckpt_dir) + "/", "--sta_dir", str(out_dir / "statistics.csv"), ] return cmd, env, repo if model_name == "cgnet": repo = ROOT / "model_repos" / "CGNet-CD" if args.dry_run: train_view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "train" val_view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "val" else: train_view = _prepare_matched_split_view(dataset_cfg, "train") val_view = _prepare_matched_split_view(dataset_cfg, "val") cmd = [ sys.executable, "train_CGNet.py", "--epoch", str(int(model_cfg.get("num_epochs", 50))), "--batchsize", str(int(dataset_cfg.get("batch_size", 8))), "--trainsize", str(int(dataset_cfg.get("img_size", 256))), "--gpu_id", child_gpu, "--data_name", dataset_cfg["name"], "--model_name", "CGNet", "--save_path", str(out_dir) + "/", "--train_root", str(train_view) + "/", "--val_root", str(val_view) + "/", ] return cmd, env, repo if model_name == "hanet": repo = ROOT / "model_repos" / "HANet-CD" view_root = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "hanet_matched" if not args.dry_run: for split in ("train", "val"): split_view = _prepare_matched_split_view(dataset_cfg, split) _safe_link_or_copy(split_view, view_root / split) env["HANET_METADATA_JSON"] = str(_write_hanet_metadata(dataset_cfg, model_cfg, view_root)) return [sys.executable, "trainHCX.py"], env, repo if model_name == "tinycd": repo = ROOT / "model_repos" / "Tiny_model_4_CD" view = ROOT / "generated_dataset_views" / dataset_cfg["name"] / "tinycd" if args.dry_run else _prepare_tinycd_view(dataset_cfg) cmd = [ sys.executable, "training.py", "--datapath", str(view), "--log-path", str(out_dir / "logs"), "--gpu-id", child_gpu, "--batch-size", str(int(args.batch_size or dataset_cfg.get("batch_size", 8))), "--epochs", str(int(args.epochs or model_cfg.get("num_epochs", 100))), ] return cmd, env, repo if model_name == "changer": repo = ROOT / "model_repos" / "open-cd" view = ( ROOT / "generated_dataset_views" / dataset_cfg["name"] / "opencd" if args.dry_run else _prepare_opencd_view(dataset_cfg) ) cfg_path = _write_opencd_changer_config(dataset_cfg, model_cfg, view) return [sys.executable, "tools/train.py", str(cfg_path), "--work-dir", str(out_dir / "work_dir")], env, repo if model_name == "changemamba": repo = ROOT / "model_repos" / "ChangeMamba" changedetection = repo / "changedetection" view = ( ROOT / "generated_dataset_views" / dataset_cfg["name"] / "changemamba" if args.dry_run else _prepare_changemamba_view(dataset_cfg) ) variant = str(model_cfg.get("vmamba_variant", "small")).lower() model_type = str(model_cfg.get("changemamba_model_type", f"MambaBCD_{variant.capitalize()}")) pretrained_name = { "tiny": "vssmtiny_dp01_ckpt_epoch_292.pth", "small": "vssmsmall_dp03_ckpt_epoch_238.pth", "base": "vssmbase_dp06_ckpt_epoch_241.pth", }.get(variant) cfg_name = { "tiny": "vssm_tiny_224_0229flex.yaml", "small": "vssm_small_224.yaml", "base": "vssm_base_224.yaml", }.get(variant) if pretrained_name is None or cfg_name is None: raise ValueError(f"Unsupported ChangeMamba VMamba variant: {variant}") pretrained = repo / "pretrained_weight" / pretrained_name cfg_path = changedetection / "configs" / "vssm1" / cfg_name max_iters, max_iters_source = _iteration_budget("changemamba", model_cfg, args, default=50000) print(f"[TRAIN-CFG] smoke_test={bool(args.smoke_test)}") print(f"[TRAIN-CFG] max_iters={max_iters}") print(f"[TRAIN-CFG] max_iters_source={max_iters_source}") cmd = [ sys.executable, "script/train_MambaBCD.py", "--dataset", _legacy_dataset_token(dataset_cfg), "--batch_size", str(int(dataset_cfg.get("batch_size", 8))), "--num_workers", str(int(dataset_cfg.get("num_workers", 4))), "--crop_size", str(int(dataset_cfg.get("img_size", model_cfg.get("img_size", 256)))), "--max_iters", str(max_iters), "--model_type", model_type, "--model_param_path", str(out_dir / "checkpoints"), "--train_dataset_path", str(view / "train"), "--train_data_list_path", str(view / "train_set.txt"), "--test_dataset_path", str(view / "test"), "--test_data_list_path", str(view / "test_set.txt"), "--cfg", str(cfg_path), "--encoder_pretrained_path", str(pretrained), "--learning_rate", str(float(model_cfg.get("lr", 6e-5))), "--weight_decay", str(float(model_cfg.get("weight_decay", 0.01))), ] return cmd, env, changedetection if model_name in {"fc_ef", "fc_siam_conc", "fc_siam_diff"}: cmd = [ sys.executable, str(ROOT / "train" / "fc_adapter.py"), "--model", model_name, "--dataset", dataset_cfg["name"], "--gpu", str(args.gpu), ] if args.epochs is not None: cmd.extend(["--epochs", str(args.epochs)]) if args.batch_size is not None: cmd.extend(["--batch-size", str(args.batch_size)]) if args.lr is not None: cmd.extend(["--lr", str(args.lr)]) if args.resume: cmd.append("--resume") if args.eval_only: cmd.append("--eval-only") if args.force: cmd.append("--force") if args.output_dir: cmd.extend(["--output-dir", args.output_dir]) return cmd, env, ROOT if model_name == "elgcnet": repo = ROOT / "model_repos" / "elgcnet" view = ( ROOT / "generated_dataset_views" / dataset_cfg["name"] / "legacy_list" if args.dry_run else prepare_legacy_list_view(dataset_cfg) ) env["CD_MODELS_DATASET_ROOT"] = str(view) env["CD_MODELS_DSIFN_ROOT"] = str(view) cmd = [ sys.executable, "main_cd.py", "--gpu_ids", child_gpu, "--project_name", f"{dataset_cfg['name']}-train-{model_name}", "--data_name", _legacy_dataset_token(dataset_cfg), "--img_size", str(int(dataset_cfg.get("img_size", model_cfg.get("img_size", 256)))), "--batch_size", str(int(dataset_cfg.get("batch_size", 8))), "--num_workers", str(int(dataset_cfg.get("num_workers", 4))), "--max_epochs", str(int(model_cfg.get("num_epochs", 200))), "--optimizer", str(model_cfg.get("optimizer", "adamw")), "--lr", str(float(model_cfg.get("lr", 0.00031))), "--loss", "ce", "--net_G", "ELGCNet", ] return cmd, env, repo return [], env, ROOT def main_for(model_name: str) -> int: args = parse_args(model_name) gpu_resolution = resolve_gpu(args.gpu) print_gpu_diagnostics(gpu_resolution) selected = args.model if model_name == "fc_variants" else model_name if selected == "fc_variants": print("[SKIP] fc_variants requires --model fc_ef, --model fc_siam_conc, or --model fc_siam_diff.") return 0 if args.dry_run else SKIP_EXIT_CODE if args.smoke_test: return run_smoke_test(selected, args) if selected in EXISTING_LEGACY: return run_existing_model(selected, args) return run_external_model(selected, args)