from __future__ import annotations import importlib import sys from abc import ABC, abstractmethod from contextlib import contextmanager from dataclasses import dataclass from pathlib import Path from types import SimpleNamespace from typing import Any import torch import torch.nn as nn import torch.nn.functional as F from utils.metrics import normalize_binary_prediction ROOT = Path(__file__).resolve().parents[1] @dataclass class NormalizedOutput: binary: torch.Tensor score: torch.Tensor | None metric_tensor: torch.Tensor class AdapterUnavailable(RuntimeError): pass @contextmanager def repo_import_context(repo: Path): repo_str = str(repo) old_path = list(sys.path) upstream_prefixes = ("models", "model", "network", "cd", "DSIFN", "changedetection") for name in list(sys.modules): if any(name == prefix or name.startswith(f"{prefix}.") for prefix in upstream_prefixes): del sys.modules[name] sys.path.insert(0, repo_str) try: yield finally: sys.path = old_path class BaseModelAdapter(ABC): model_name = "base" supports_inprocess_eval = False supports_unified_eval = False supports_unified_training = False supports_flops = False supports_amp_training = True notes_or_failure_reason = "Adapter is not implemented." model_class_path = "" input_format = "" output_format = "" checkpoint_format = "" final_output_for_metrics = "" def __init__(self, model_name: str | None = None) -> None: if model_name: self.model_name = model_name def _missing(self, method: str) -> NotImplementedError: return NotImplementedError(f"{self.model_name} adapter does not implement {method}: {self.notes_or_failure_reason}") def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: raise self._missing("build_model") def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: raise self._missing("load_checkpoint") def forward(self, model: nn.Module, batch: tuple, device: torch.device) -> Any: raise self._missing("forward") def normalize_output(self, raw_output: Any, batch: tuple, dataset_config: dict) -> NormalizedOutput: raise self._missing("normalize_output") def get_dummy_inputs(self, dataset_config: dict, device: torch.device) -> tuple: raise self._missing("get_dummy_inputs") def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: output = raw_output[-1] if isinstance(raw_output, (list, tuple)) else raw_output output = output.float() mask = batch[2].to(device, non_blocking=True).float() if output.ndim == 4 and output.shape[1] == 2: target = mask.squeeze(1).long() loss = F.cross_entropy(output, target) elif output.ndim == 4 and output.shape[1] == 1: target = mask if output.shape[-2:] != target.shape[-2:]: output = F.interpolate(output, size=target.shape[-2:], mode="bilinear", align_corners=True) if output.detach().min() >= 0 and output.detach().max() <= 1: loss = F.binary_cross_entropy(output.clamp(1e-6, 1 - 1e-6), target) else: loss = F.binary_cross_entropy_with_logits(output, target) else: raise RuntimeError(f"{self.model_name} default loss cannot handle output shape {tuple(output.shape)}.") return {"loss": loss, "main_loss": loss.detach()} def save_checkpoint( self, model: nn.Module, optimizer: torch.optim.Optimizer, scheduler: Any, path: Path, metadata: dict, ) -> None: payload = dict(metadata) payload["model_state_dict"] = model.state_dict() payload["optimizer_state_dict"] = optimizer.state_dict() if scheduler is not None: payload["scheduler_state_dict"] = scheduler.state_dict() torch.save(payload, path) def build_optimizer(self, model: nn.Module, model_config: dict) -> torch.optim.Optimizer: lr = float(model_config.get("lr", 1e-4)) weight_decay = float(model_config.get("weight_decay", 0.0) or 0.0) optimizer_name = str(model_config.get("optimizer", "adam")).lower() if optimizer_name == "sgd": return torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=weight_decay) if optimizer_name == "adamw": return torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay) return torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay) def build_scheduler(self, optimizer: torch.optim.Optimizer, model_config: dict) -> Any: scheduler_name = str(model_config.get("scheduler", "")).lower() epochs = int(model_config.get("num_epochs", 200)) if scheduler_name == "cosine": return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) if scheduler_name == "step": return torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5) return None def get_primary_metric(self) -> str: return "f1" def get_threshold(self, dataset_config: dict) -> float: return float(dataset_config.get("eval", {}).get("threshold", 0.5)) def get_threshold_mode(self) -> str: return "argmax" if "two-channel" in self.output_format else "threshold" class UnsupportedAdapter(BaseModelAdapter): def __init__(self, model_name: str, reason: str) -> None: super().__init__(model_name) self.notes_or_failure_reason = reason class TwoTensorLogitAdapter(BaseModelAdapter): supports_inprocess_eval = True supports_unified_eval = True supports_unified_training = True supports_flops = True input_format = "two tensors (A, B), each [B, 3, H, W]" output_format = "two-channel logits [B, 2, H, W]" final_output_for_metrics = "raw two-channel logits" def forward(self, model: nn.Module, batch: tuple, device: torch.device) -> torch.Tensor: a, b, _mask, _names = batch return model(a.to(device, non_blocking=True), b.to(device, non_blocking=True)) def normalize_output(self, raw_output: Any, batch: tuple, dataset_config: dict) -> NormalizedOutput: if isinstance(raw_output, (list, tuple)): raw_output = raw_output[-1] threshold = float(dataset_config.get("eval", {}).get("threshold", 0.5)) binary, score = normalize_binary_prediction(raw_output.detach().cpu(), threshold=threshold) return NormalizedOutput(binary=binary, score=score, metric_tensor=raw_output.detach().cpu()) def get_dummy_inputs(self, dataset_config: dict, device: torch.device) -> tuple: size = int(dataset_config.get("img_size", dataset_config.get("image_size", 256))) return ( torch.zeros(1, 3, size, size, device=device), torch.zeros(1, 3, size, size, device=device), ) class FCAdapter(TwoTensorLogitAdapter): notes_or_failure_reason = "Direct FC model from fully_convolutional_change_detection." checkpoint_format = "CD-Models checkpoint dict with model_state_dict, or raw state dict" def __init__(self, model_name: str) -> None: super().__init__(model_name) self.model_class_path = { "fc_ef": "model_repos/fully_convolutional_change_detection/unet.py:Unet", "fc_siam_conc": "model_repos/fully_convolutional_change_detection/siamunet_conc.py:SiamUnet_conc", "fc_siam_diff": "model_repos/fully_convolutional_change_detection/siamunet_diff.py:SiamUnet_diff", }[model_name] def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "model_repos" / "fully_convolutional_change_detection" with repo_import_context(repo): if self.model_name == "fc_ef": module = importlib.import_module("unet") model = module.Unet(input_nbr=6, label_nbr=2) elif self.model_name == "fc_siam_conc": module = importlib.import_module("siamunet_conc") model = module.SiamUnet_conc(input_nbr=3, label_nbr=2) else: module = importlib.import_module("siamunet_diff") model = module.SiamUnet_diff(input_nbr=3, label_nbr=2) return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"{self.model_name} checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: target = batch[2].squeeze(1).long().to(device, non_blocking=True) weight = None if self.model_name == "fc_siam_diff": pos = target.eq(1).sum().float() neg = target.eq(0).sum().float() if bool(pos.item() > 0): pos_weight = (neg / pos.clamp_min(1.0)).clamp(min=1.0, max=50.0) weight = torch.stack([torch.ones_like(pos_weight), pos_weight]).to(device) loss = F.nll_loss(raw_output.float(), target, weight=weight) result = {"loss": loss, "nll_loss": loss.detach()} if weight is not None: result["positive_class_weight"] = weight[1].detach() return result class BITAdapter(TwoTensorLogitAdapter): model_name = "bit_cd" notes_or_failure_reason = "BIT_CD models.networks.define_G with base_transformer_pos_s4_dd8." model_class_path = "BIT_CD/models/networks.py:define_G" checkpoint_format = "raw state dict, CD-Models model_state_dict, or upstream model_G_state_dict" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "BIT_CD" with repo_import_context(repo): networks = importlib.import_module("models.networks") args = SimpleNamespace(net_G="base_transformer_pos_s4_dd8", lr_policy="linear", max_epochs=1) model = networks.define_G(args, gpu_ids=[]) return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_G_state_dict" in checkpoint: state = checkpoint["model_G_state_dict"] elif isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"BIT_CD checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) class ChangeFormerAdapter(BITAdapter): model_name = "changeformer" supports_amp_training = False notes_or_failure_reason = "ChangeFormer models.networks.define_G with ChangeFormerV6 embed_dim=64." model_class_path = "ChangeFormer/models/networks.py:define_G" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "ChangeFormer" with repo_import_context(repo): networks = importlib.import_module("models.networks") args = SimpleNamespace(net_G="ChangeFormerV6", embed_dim=64, lr_policy="linear", max_epochs=1) model = networks.define_G(args, gpu_ids=[]) return model.to(device) class BiFAAdapter(TwoTensorLogitAdapter): model_name = "bifa" supports_amp_training = False notes_or_failure_reason = "BiFA models.bifa.BiFA with mit_b0 backbone; returns two-channel logits." model_class_path = "BiFA/models/bifa.py:BiFA" output_format = "two-channel logits [B, 2, H, W]" checkpoint_format = "raw state dict or CD-Models model_state_dict" final_output_for_metrics = "raw two-channel logits" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "BiFA" with repo_import_context(repo): module = importlib.import_module("models.bifa") model = module.BiFA(backbone="mit_b0") return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"BiFA checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) class SiamNestedUNetAdapter(TwoTensorLogitAdapter): model_name = "siam_nestedunet" supports_amp_training = False notes_or_failure_reason = "Siam-NestedUNet SNUNet_ECAM returns a tuple; final output is tuple[-1]." model_class_path = "Siam-NestedUNet/models/Models.py:SNUNet_ECAM" output_format = "tuple containing two-channel logits [B, 2, H, W]" checkpoint_format = "raw state dict or CD-Models model_state_dict" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "Siam-NestedUNet" with repo_import_context(repo): module = importlib.import_module("models.Models") model = module.SNUNet_ECAM(3, 2) return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"Siam-NestedUNet checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) class STANetModule(nn.Module): def __init__(self, net_f: nn.Module, net_a: nn.Module) -> None: super().__init__() self.netF = net_f self.netA = net_a def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: feat_a = self.netF(a) feat_b = self.netF(b) feat_a, feat_b = self.netA(feat_a, feat_b) dist = torch.norm(feat_a - feat_b, p=2, dim=1, keepdim=True) return F.interpolate(dist, size=a.shape[-2:], mode="bilinear", align_corners=True) class STANetAdapter(BaseModelAdapter): model_name = "stanet" supports_inprocess_eval = True supports_unified_eval = True supports_unified_training = True supports_flops = True notes_or_failure_reason = "STANet netF + CDSA(PAM), distance map thresholded at 1.0 as in train_wildfire.py." model_class_path = "STANet/models/backbone.py:define_F + CDSA" input_format = "two tensors (A, B), each [B, 3, H, W]" output_format = "one-channel L2 distance map [B, 1, H, W]" checkpoint_format = "dict with netF and netA state dicts" final_output_for_metrics = "binary distance > 1.0" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "STANet" with repo_import_context(repo): backbone = importlib.import_module("models.backbone") net_f = backbone.define_F(in_c=3, f_c=64, type="mynet3") net_a = backbone.CDSA(in_c=64, ds=1, mode="PAM") return STANetModule(net_f, net_a).to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: model.load_state_dict(checkpoint["model_state_dict"], strict=True) return if not isinstance(checkpoint, dict) or "netF" not in checkpoint or "netA" not in checkpoint: raise RuntimeError(f"STANet checkpoint {checkpoint_path} must contain netF and netA keys.") model.netF.load_state_dict(checkpoint["netF"], strict=True) model.netA.load_state_dict(checkpoint["netA"], strict=True) def forward(self, model: nn.Module, batch: tuple, device: torch.device) -> torch.Tensor: a, b, _mask, _names = batch return model(a.to(device, non_blocking=True), b.to(device, non_blocking=True)) def normalize_output(self, raw_output: Any, batch: tuple, dataset_config: dict) -> NormalizedOutput: dist = raw_output.detach().cpu() binary = (dist > 1.0).float() return NormalizedOutput(binary=binary, score=None, metric_tensor=binary) def get_dummy_inputs(self, dataset_config: dict, device: torch.device) -> tuple: size = int(dataset_config.get("img_size", dataset_config.get("image_size", 256))) return ( torch.zeros(1, 3, size, size, device=device), torch.zeros(1, 3, size, size, device=device), ) def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: repo = ROOT / "STANet" with repo_import_context(repo): module = importlib.import_module("models.loss") criterion = module.BCL().to(device) target = batch[2].to(device, non_blocking=True).float() label = target.clone() label[target == 1] = -1 label[target == 0] = 1 loss = criterion(raw_output, label) return {"loss": loss, "bcl_loss": loss.detach()} class TinyCDAdapter(TwoTensorLogitAdapter): model_name = "tinycd" supports_inprocess_eval = True supports_unified_eval = True supports_unified_training = True supports_flops = True supports_amp_training = False notes_or_failure_reason = "TinyCD ChangeClassifier, sigmoid one-channel probability output." model_class_path = "model_repos/Tiny_model_4_CD/models/change_classifier.py:ChangeClassifier" output_format = "one-channel sigmoid probability map [B, 1, H, W]" checkpoint_format = "raw state dict or CD-Models model_state_dict" final_output_for_metrics = "sigmoid probability thresholded with dataset eval threshold" def get_threshold_mode(self) -> str: return "threshold" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "model_repos" / "Tiny_model_4_CD" with repo_import_context(repo): module = importlib.import_module("models.change_classifier") model = module.ChangeClassifier() return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"TinyCD checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: target = batch[2].to(device, non_blocking=True).float() output = raw_output.float() if output.shape[-2:] != target.shape[-2:]: output = F.interpolate(output, size=target.shape[-2:], mode="bilinear", align_corners=True) loss = F.binary_cross_entropy(output.clamp(1e-6, 1 - 1e-6), target) return {"loss": loss, "bce_loss": loss.detach()} class DSAMNetAdapter(BaseModelAdapter): model_name = "dsamnet" supports_inprocess_eval = True supports_unified_eval = True supports_unified_training = True supports_flops = True notes_or_failure_reason = "DSAMNet returns distance map plus deep supervision outputs; distance > 1 is change." model_class_path = "model_repos/DSAMNet/model/dsamnet.py:DSAMNet" input_format = "two tensors (A, B), each [B, 3, H, W]" output_format = "tuple(dist [B,1,H,W], ds2, ds3)" checkpoint_format = "raw state dict or CD-Models model_state_dict" final_output_for_metrics = "dist > 1.0" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "model_repos" / "DSAMNet" with repo_import_context(repo): module = importlib.import_module("model.dsamnet") model = module.DSAMNet(2) return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"DSAMNet checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) def forward(self, model: nn.Module, batch: tuple, device: torch.device) -> Any: a, b, _mask, _names = batch return model(a.to(device, non_blocking=True), b.to(device, non_blocking=True)) def normalize_output(self, raw_output: Any, batch: tuple, dataset_config: dict) -> NormalizedOutput: dist = raw_output[0].detach().cpu() if isinstance(raw_output, (list, tuple)) else raw_output.detach().cpu() binary = (dist > 1.0).float() return NormalizedOutput(binary=binary, score=None, metric_tensor=binary) def get_dummy_inputs(self, dataset_config: dict, device: torch.device) -> tuple: size = int(dataset_config.get("img_size", dataset_config.get("image_size", 256))) return ( torch.zeros(1, 3, size, size, device=device), torch.zeros(1, 3, size, size, device=device), ) def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: if not isinstance(raw_output, (list, tuple)) or len(raw_output) < 3: raise RuntimeError("DSAMNet training expects (dist, ds2, ds3).") dist, ds2, ds3 = raw_output[:3] repo = ROOT / "model_repos" / "DSAMNet" with repo_import_context(repo): bcl_module = importlib.import_module("loss.BCL") dice_module = importlib.import_module("loss.DiceLoss") bcl = bcl_module.BCL().to(device) dice = dice_module.DiceLoss().to(device) mask = batch[2].to(device, non_blocking=True).float() one_hot = torch.cat([1.0 - mask, mask], dim=1) dice_loss = 0.5 * (dice(ds2, one_hot) + dice(ds3, one_hot)) ct_loss = bcl(dist, mask.clone()) total = ct_loss + float(model_config.get("wDice", 0.5)) * dice_loss return {"loss": total, "contrastive_loss": ct_loss.detach(), "dice_loss": dice_loss.detach()} class CGNetAdapter(TwoTensorLogitAdapter): model_name = "cgnet" notes_or_failure_reason = "CGNet returns (change_map, final_map); final_map is used for metrics." model_class_path = "model_repos/CGNet-CD/network/CGNet.py:CGNet" output_format = "tuple(one-channel coarse logits, one-channel final logits)" checkpoint_format = "raw state dict or CD-Models model_state_dict" final_output_for_metrics = "tuple[-1] final_map" def get_threshold_mode(self) -> str: return "threshold" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "model_repos" / "CGNet-CD" with repo_import_context(repo): module = importlib.import_module("network.CGNet") model = module.CGNet() return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"CGNet checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: if not isinstance(raw_output, (list, tuple)) or len(raw_output) < 2: raise RuntimeError("CGNet training expects a tuple of two prediction maps.") target = batch[2].to(device, non_blocking=True).float() loss1 = F.binary_cross_entropy_with_logits(raw_output[0], target) loss2 = F.binary_cross_entropy_with_logits(raw_output[1], target) loss = loss1 + loss2 return {"loss": loss, "coarse_loss": loss1.detach(), "final_loss": loss2.detach()} class HANetAdapter(TwoTensorLogitAdapter): model_name = "hanet" supports_amp_training = False notes_or_failure_reason = "HANet HAN returns a tuple; tuple[-1] is used for metrics." model_class_path = "model_repos/HANet-CD/models/HANet.py:HAN" output_format = "tuple containing two-channel logits [B, 2, H, W]" checkpoint_format = "full torch-saved model, raw state dict, or CD-Models model_state_dict" final_output_for_metrics = "tuple[-1]" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "model_repos" / "HANet-CD" with repo_import_context(repo): module = importlib.import_module("models.HANet") model = module.HAN(3, 2) return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, nn.Module): model.load_state_dict(checkpoint.state_dict(), strict=True) return if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"HANet checkpoint {checkpoint_path} has no recognized model/state dict.") model.load_state_dict(state, strict=True) def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: logits = raw_output[-1] if isinstance(raw_output, (list, tuple)) else raw_output logits = logits.float() target = batch[2].long().to(device, non_blocking=True) if target.ndim == 4: target = target.squeeze(1) ce = F.cross_entropy(logits, target) probas = F.softmax(logits, dim=1) one_hot = F.one_hot(target, num_classes=logits.shape[1]).permute(0, 3, 1, 2).float() dims = (0, 2, 3) dice = 1.0 - ((2.0 * torch.sum(probas * one_hot, dims) + 1e-7) / (torch.sum(probas + one_hot, dims) + 1e-7)).mean() loss = ce + dice return {"loss": loss, "cross_entropy": ce.detach(), "dice_loss": dice.detach()} class ELGCNetAdapter(TwoTensorLogitAdapter): model_name = "elgcnet" notes_or_failure_reason = "ELGCNet models.networks.define_G with net_G=ELGCNet and dec_embed_dim=256." model_class_path = "model_repos/elgcnet/models/networks.py:define_G" output_format = "two-channel logits [B, 2, H, W]" checkpoint_format = "upstream model_G_state_dict, raw state dict, or CD-Models model_state_dict" final_output_for_metrics = "raw two-channel logits" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "model_repos" / "elgcnet" with repo_import_context(repo): networks = importlib.import_module("models.networks") args = SimpleNamespace(net_G="ELGCNet", dec_embed_dim=256) model = networks.define_G(args, gpu_ids=[]) return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_G_state_dict" in checkpoint: state = checkpoint["model_G_state_dict"] elif isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"ELGCNet checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) class DSIFNAdapter(BaseModelAdapter): supports_inprocess_eval = True supports_unified_eval = True supports_unified_training = True supports_flops = True supports_amp_training = False notes_or_failure_reason = "IFNet/pytorch version DSIFN with two VGG16 feature towers; upstream validation uses preds[-1]." model_class_path = "IFNet/pytorch version/DSIFN.py:DSIFN" input_format = "two tensors (A, B), each [B, 3, H, W]" output_format = "tuple of five one-channel sigmoid maps" checkpoint_format = "raw state dict, CD-Models model_state_dict, or upstream state_dict" final_output_for_metrics = "tuple[-1], resized to target mask then thresholded" def __init__(self, model_name: str) -> None: super().__init__(model_name) def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "IFNet" / "pytorch version" with repo_import_context(repo): import torchvision.models as tv_models original_vgg16 = tv_models.vgg16 def vgg16_no_download(*args: Any, **kwargs: Any) -> nn.Module: kwargs.pop("pretrained", None) kwargs["weights"] = None return original_vgg16(*args, **kwargs) tv_models.vgg16 = vgg16_no_download try: module = importlib.import_module("DSIFN") model = module.DSIFN(module.vgg16_base(), module.vgg16_base()) finally: tv_models.vgg16 = original_vgg16 return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] if any(str(key).startswith("net.") for key in state): model.load_state_dict(state, strict=True) return elif isinstance(checkpoint, dict) and "state_dict" in checkpoint: state = checkpoint["state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"{self.model_name} checkpoint {checkpoint_path} has no recognized state dict.") model.load_state_dict(state, strict=True) def forward(self, model: nn.Module, batch: tuple, device: torch.device) -> Any: a, b, _mask, _names = batch return model(a.to(device, non_blocking=True), b.to(device, non_blocking=True)) def normalize_output(self, raw_output: Any, batch: tuple, dataset_config: dict) -> NormalizedOutput: if not isinstance(raw_output, (list, tuple)) or not raw_output: raise RuntimeError(f"{self.model_name} expected DSIFN tuple output, got {type(raw_output)!r}.") score = raw_output[-1].detach().cpu() target = batch[2] if score.shape[-2:] != target.shape[-2:]: score = F.interpolate(score, size=target.shape[-2:], mode="bilinear", align_corners=True) threshold = float(dataset_config.get("eval", {}).get("threshold", 0.5)) binary, score = normalize_binary_prediction(score, threshold=threshold) return NormalizedOutput(binary=binary, score=score, metric_tensor=score) def get_dummy_inputs(self, dataset_config: dict, device: torch.device) -> tuple: size = int(dataset_config.get("img_size", dataset_config.get("image_size", 256))) return ( torch.zeros(1, 3, size, size, device=device), torch.zeros(1, 3, size, size, device=device), ) def get_threshold_mode(self) -> str: return "threshold" def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: if not isinstance(raw_output, (list, tuple)) or not raw_output: raise RuntimeError(f"{self.model_name} training expects five DSIFN outputs.") target = batch[2].to(device, non_blocking=True).float() repo = ROOT / "IFNet" / "pytorch version" with repo_import_context(repo): loss_module = importlib.import_module("loss") cd_loss = loss_module.cd_loss total = torch.zeros((), device=device) for pred in raw_output: pred = pred.float() if pred.shape[-2:] != target.shape[-2:]: pred = F.interpolate(pred, size=target.shape[-2:], mode="bilinear", align_corners=True) total = total + cd_loss(pred.squeeze(1), target.squeeze(1)) loss = total / len(raw_output) return {"loss": loss, "dsifn_loss": loss.detach()} class SChangerModule(nn.Module): def __init__(self, net: nn.Module) -> None: super().__init__() self.net = net def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return self.net((a, b)) class SChangerAdapter(TwoTensorLogitAdapter): model_name = "schanger" notes_or_failure_reason = "SChanger-base from SChanger/cd/SChanger.py; eval forward returns sigmoid probability map." model_class_path = "SChanger/cd/SChanger.py:SChanger" output_format = "one-channel sigmoid probability map [B, 1, H, W] in eval mode" checkpoint_format = "raw state dict, CD-Models model_state_dict, or upstream state_dict" final_output_for_metrics = "eval-mode sigmoid probability thresholded with dataset eval threshold" def get_threshold_mode(self) -> str: return "threshold" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "SChanger" with repo_import_context(repo): module = importlib.import_module("cd.SChanger") net = module.SChanger( num_classes=1, input_channels=3, c_list=[8 * 3, 8 * 4, 8 * 6, 8 * 8, 8 * 13, 8 * 15], dropout=0.2, ) return SChangerModule(net).to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] if any(str(key).startswith(("net.", "trainer.")) for key in state): model.load_state_dict(state, strict=True) return elif isinstance(checkpoint, dict) and "state_dict" in checkpoint: state = checkpoint["state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"SChanger checkpoint {checkpoint_path} has no recognized state dict.") model.net.load_state_dict(state, strict=True) def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: target = batch[2].to(device, non_blocking=True).float() output = raw_output if output.ndim != 4: raise RuntimeError(f"SChanger training expected [B,C,H,W] output, got {tuple(output.shape)}.") total = torch.zeros((), device=device) for idx in range(output.shape[1]): total = total + F.binary_cross_entropy_with_logits(output[:, idx:idx + 1], target) loss = total / output.shape[1] return {"loss": loss, "deep_supervision_bce": loss.detach()} class Change3DBCDModule(nn.Module): def __init__(self, trainer: nn.Module) -> None: super().__init__() self.trainer = trainer def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return self.trainer.update_bcd(a, b) class Change3DAdapter(TwoTensorLogitAdapter): model_name = "change3d" notes_or_failure_reason = "Change3D Trainer.update_bcd with X3D-L encoder and ChangeDecoder sigmoid output." model_class_path = "Change3D/model/trainer.py:Trainer" output_format = "one-channel sigmoid probability map [B, 1, H, W]" checkpoint_format = "raw Trainer state dict, CD-Models model_state_dict, or upstream state_dict" final_output_for_metrics = "update_bcd sigmoid probability thresholded with dataset eval threshold" def get_threshold_mode(self) -> str: return "threshold" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "Change3D" size = int(dataset_config.get("img_size", dataset_config.get("image_size", 256))) pretrained = repo / "model" / "X3D_L.pyth" args = SimpleNamespace( dataset=str(dataset_config.get("name", "")).upper() + "_CD", in_height=size, in_width=size, num_perception_frame=1, num_class=1, pretrained=str(pretrained) if pretrained.exists() else None, batch_size=int(dataset_config.get("batch_size", 1)), lr=float(model_config.get("lr", 0.0002)), lr_mode="poly", max_steps=1, max_epochs=1, step_loss=100, ) with repo_import_context(repo): module = importlib.import_module("model.trainer") trainer = module.Trainer(args) return Change3DBCDModule(trainer).to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: checkpoint = torch.load(checkpoint_path, map_location=device) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: state = checkpoint["model_state_dict"] elif isinstance(checkpoint, dict) and "state_dict" in checkpoint: state = checkpoint["state_dict"] elif isinstance(checkpoint, dict) and all(torch.is_tensor(v) for v in checkpoint.values()): state = checkpoint else: raise RuntimeError(f"Change3D checkpoint {checkpoint_path} has no recognized state dict.") model.trainer.load_state_dict(state, strict=True) def compute_loss(self, raw_output: Any, batch: tuple, model_config: dict, dataset_config: dict, device: torch.device) -> dict[str, torch.Tensor]: target = batch[2].to(device, non_blocking=True).float() output = raw_output.float() if output.shape[-2:] != target.shape[-2:]: output = F.interpolate(output, size=target.shape[-2:], mode="bilinear", align_corners=True) bce = F.binary_cross_entropy(output.clamp(1e-6, 1 - 1e-6), target) smooth = 1.0 intersection = torch.sum(output * target) dice = 1.0 - ((2.0 * intersection + smooth) / (torch.sum(output) + torch.sum(target) + smooth)) loss = bce + dice return {"loss": loss, "bce_loss": bce.detach(), "dice_loss": dice.detach()} class ChangeMambaAdapter(TwoTensorLogitAdapter): model_name = "changemamba" supports_unified_training = False notes_or_failure_reason = "ChangeMambaBCD with VMamba backbone; upstream training checkpoint stores weights under the model key." model_class_path = "model_repos/ChangeMamba/changedetection/models/ChangeMambaBCD.py:ChangeMambaBCD" output_format = "two-channel logits [B, 2, H, W]" checkpoint_format = "ChangeMamba training checkpoint with model key, state_dict key, or raw state dict" final_output_for_metrics = "raw two-channel logits" def build_model(self, model_config: dict, dataset_config: dict, device: torch.device) -> nn.Module: repo = ROOT / "model_repos" / "ChangeMamba" variant = str(model_config.get("vmamba_variant", "small")).lower() cfg_name = { "tiny": "vssm_tiny_224_0229flex.yaml", "small": "vssm_small_224.yaml", "base": "vssm_base_224.yaml", }.get(variant) if cfg_name is None: raise ValueError(f"Unsupported ChangeMamba VMamba variant: {variant}") cfg_path = repo / "changedetection" / "configs" / "vssm1" / cfg_name args = SimpleNamespace( cfg=str(cfg_path), opts=None, batch_size=int(dataset_config.get("batch_size", 8)), data_path=None, zip=None, cache_mode=None, pretrained=None, encoder_pretrained_path=None, model_checkpoint_path=None, resume=None, resume_training_path=None, accumulation_steps=None, use_checkpoint=None, disable_amp=None, output=None, tag=None, eval=None, throughput=None, enable_amp=None, fused_layernorm=None, optim=None, ) with repo_import_context(repo): config_module = importlib.import_module("changedetection.configs.config") utils_module = importlib.import_module("changedetection.script.script_utils") model_module = importlib.import_module("changedetection.models.ChangeMambaBCD") config = config_module.get_config(args) model = model_module.ChangeMambaBCD(pretrained=None, **utils_module.get_vssm_kwargs(config)) return model.to(device) def load_checkpoint(self, model: nn.Module, checkpoint_path: Path, device: torch.device) -> None: # Load large upstream ChangeMamba training archives on CPU first. # Direct CUDA deserialization can fail with PytorchStreamReader read # errors on otherwise valid zip checkpoints, while load_state_dict will # copy the extracted tensors into the already-placed CUDA model. checkpoint = torch.load(checkpoint_path, map_location="cpu") with repo_import_context(ROOT / "model_repos" / "ChangeMamba"): checkpoint_module = importlib.import_module("changedetection.checkpoints") state = checkpoint_module.extract_model_state_dict(checkpoint) model.load_state_dict(state, strict=True) _SUPPORTED = { "fc_ef": FCAdapter, "fc_siam_conc": FCAdapter, "fc_siam_diff": FCAdapter, "bifa": BiFAAdapter, "bit_cd": BITAdapter, "changeformer": ChangeFormerAdapter, "changemamba": ChangeMambaAdapter, "change3d": Change3DAdapter, "siam_nestedunet": SiamNestedUNetAdapter, "stanet": STANetAdapter, "tinycd": TinyCDAdapter, "dsamnet": DSAMNetAdapter, "dsifn": DSIFNAdapter, "ifnet": DSIFNAdapter, "cgnet": CGNetAdapter, "hanet": HANetAdapter, "elgcnet": ELGCNetAdapter, "schanger": SChangerAdapter, } _UNSUPPORTED_REASONS = { "cdmamba": "CDMamba requires the cd-mamba-ssm environment with CUDA-visible mamba_ssm/Triton; cd-mamba-ssm is not available in this shell, and Mamba reports torch.cuda.is_available() == False.", "changer": "Changer command construction is wired through Open-CD/OpenMMLab; verify runtime dependencies in cd-openmmlab before treating it as trainable.", "rsm_cd": "RSM-CD requires the cd-mamba-ssm environment and VMamba/selective-scan CUDA extensions; cd-mamba-ssm is not available in this shell.", } def get_model_adapter(model_name: str) -> BaseModelAdapter: if model_name in _SUPPORTED: cls = _SUPPORTED[model_name] return cls(model_name) if cls in {FCAdapter, DSIFNAdapter} else cls() return UnsupportedAdapter(model_name, _UNSUPPORTED_REASONS.get(model_name, "No adapter registered."))