Upload DOFA model package

.gitattributes

@@ -1,35 +1,13 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
 *.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,53 @@
+seed: 42
+data:
+  root: data
+  protocol: dofa_multisensor_npz_v1
+  image_size: 224
+  data_range: 1.0
+  train_samples_per_modality: 1
+  test_samples_per_modality: 1
+  modalities:
+    sentinel1:
+      channels: 2
+      wavelengths: [3.5, 4.0]
+    sentinel2:
+      channels: 9
+      wavelengths: [0.443, 0.490, 0.560, 0.665, 0.705, 0.740, 0.783, 0.842, 1.610]
+    naip:
+      channels: 3
+      wavelengths: [0.490, 0.560, 0.665]
+    gaofen:
+      channels: 4
+      wavelengths: [0.490, 0.560, 0.665, 0.830]
+    enmap:
+      channels: 202
+      wavelength_mode: synthetic_uniform
+      wavelength_start: 0.420
+      wavelength_step: 0.006
+model:
+  image_size: 224
+  patch_size: 16
+  embed_dim: 32
+  encoder_depth: 1
+  encoder_heads: 4
+  decoder_dim: 24
+  decoder_depth: 1
+  decoder_heads: 4
+  wavelength_embed_dim: 16
+  hypernetwork_heads: 4
+  mask_ratio: 0.75
+training:
+  epochs: 1
+  batch_size: 1
+  learning_rate: 0.0001
+  weight_decay: 0.05
+  num_workers: 0
+  amp: true
+runtime:
+  device: auto
+  inference_batch_size: 1
+paths:
+  checkpoint: result/checkpoints/dofa.pt
+  training_metrics: result/training/metrics.json
+  inference_dir: result/output
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,104 @@
+{
+  "model_name": "DOFA",
+  "model_type": "dofa",
+  "architectures": [
+    "DOFA"
+  ],
+  "framework": "PyTorch",
+  "domain": "earth-observation",
+  "task": "multisensor-masked-image-modeling",
+  "implementation": {
+    "entry_point": "model/dofa.py",
+    "scope": "wavelength-conditioned dynamic patch embedding and reconstruction for variable-channel Earth-observation imagery",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "family": "wavelength-conditioned masked autoencoder",
+    "image_size": 224,
+    "patch_size": 16,
+    "num_patches": 196,
+    "embed_dim": 32,
+    "encoder_depth": 1,
+    "encoder_heads": 4,
+    "decoder_dim": 24,
+    "decoder_depth": 1,
+    "decoder_heads": 4,
+    "wavelength_embed_dim": 16,
+    "hypernetwork_heads": 4,
+    "mask_ratio": 0.75,
+    "dynamic_components": [
+      "encoder patch weights",
+      "encoder patch bias",
+      "decoder reconstruction weights",
+      "decoder reconstruction bias"
+    ]
+  },
+  "data": {
+    "datasets": [
+      "Sentinel-1",
+      "Sentinel-2",
+      "NAIP",
+      "EnMAP",
+      "Gaofen"
+    ],
+    "protocol": "dofa_multisensor_npz_v1",
+    "format": "NPZ",
+    "image_key": "images",
+    "image_shape": [
+      "N",
+      "C",
+      224,
+      224
+    ],
+    "wavelength_key": "wavelengths",
+    "wavelength_shape": [
+      "C"
+    ],
+    "required_metadata": [
+      "modality",
+      "protocol"
+    ],
+    "provenance_and_metric_metadata": [
+      "data_source",
+      "data_range",
+      "wavelength_mode"
+    ],
+    "modalities": {
+      "sentinel1": {
+        "channels": 2,
+        "wavelengths": [3.5, 4.0]
+      },
+      "sentinel2": {
+        "channels": 9,
+        "wavelengths": [0.443, 0.49, 0.56, 0.665, 0.705, 0.74, 0.783, 0.842, 1.61]
+      },
+      "naip": {
+        "channels": 3,
+        "wavelengths": [0.49, 0.56, 0.665]
+      },
+      "gaofen": {
+        "channels": 4,
+        "wavelengths": [0.49, 0.56, 0.665, 0.83]
+      },
+      "enmap": {
+        "channels": 202,
+        "wavelength_mode": "synthetic_uniform",
+        "wavelength_start": 0.42,
+        "wavelength_step": 0.006
+      }
+    },
+    "default_data_range": 1.0,
+    "default_file_pattern": "data/{split}_{modality}.npz"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/dofa.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,13 @@
+{
+  "framework": "PyTorch",
+  "task": "multimodal_earth_observation_pretraining",
+  "model": "DOFA",
+  "license": "MIT",
+  "input_format": "BCHW_WITH_WAVELENGTHS",
+  "protocol": "dofa_multisensor_npz_v1",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/dofa.py

@@ -0,0 +1,134 @@
+"""Wavelength-conditioned, visible-token DOFA masked autoencoder."""
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class WavelengthEncoder(nn.Module):
+    def __init__(self, dimension):
+        super().__init__()
+        self.dimension = dimension
+        self.mlp = nn.Sequential(nn.Linear(dimension, dimension), nn.GELU(),
+                                 nn.Linear(dimension, dimension))
+
+    def forward(self, wavelengths):
+        half = self.dimension // 2
+        scale = torch.exp(torch.arange(half, device=wavelengths.device,
+                                      dtype=wavelengths.dtype) *
+                          (-math.log(10000.0) / max(half - 1, 1)))
+        phase = wavelengths[:, None] * 1000.0 * scale[None]
+        encoded = torch.cat((phase.sin(), phase.cos()), -1)
+        encoded = F.pad(encoded, (0, self.dimension - encoded.shape[-1]))
+        return encoded + self.mlp(encoded)
+
+
+class DynamicPatchWeights(nn.Module):
+    def __init__(self, wavelength_dim, embed_dim, decoder_dim, patch_size, heads):
+        super().__init__()
+        layer = nn.TransformerEncoderLayer(wavelength_dim, heads, wavelength_dim * 2,
+                                           batch_first=True, norm_first=True, dropout=0.0,
+                                           activation="gelu")
+        self.context = nn.TransformerEncoder(layer, 1)
+        self.patch_size = patch_size
+        self.embed_dim = embed_dim
+        self.decoder_dim = decoder_dim
+        self.encoder_weight = nn.Linear(wavelength_dim, embed_dim * patch_size**2)
+        self.encoder_bias = nn.Linear(wavelength_dim, embed_dim)
+        self.decoder_weight = nn.Linear(wavelength_dim, decoder_dim * patch_size**2)
+        self.decoder_bias = nn.Linear(wavelength_dim, patch_size**2)
+
+    def forward(self, wavelength_features):
+        context = self.context(wavelength_features.unsqueeze(0)).squeeze(0)
+        channels = len(context)
+        encoder_weight = self.encoder_weight(context).reshape(
+            channels, self.embed_dim, self.patch_size, self.patch_size).permute(1, 0, 2, 3)
+        encoder_bias = self.encoder_bias(context).mean(0)
+        decoder_weight = self.decoder_weight(context).reshape(
+            channels, self.patch_size**2, self.decoder_dim)
+        decoder_bias = self.decoder_bias(context)
+        return encoder_weight, encoder_bias, decoder_weight, decoder_bias
+
+
+class DOFA(nn.Module):
+    def __init__(self, image_size=224, patch_size=16, embed_dim=64, encoder_depth=2,
+                 encoder_heads=4, decoder_dim=32, decoder_depth=1, decoder_heads=4,
+                 wavelength_embed_dim=32, hypernetwork_heads=4, mask_ratio=0.75):
+        super().__init__()
+        if image_size % patch_size:
+            raise ValueError("image_size must be divisible by patch_size")
+        self.image_size, self.patch_size = image_size, patch_size
+        self.mask_ratio = mask_ratio
+        self.num_patches = (image_size // patch_size) ** 2
+        self.wavelength_encoder = WavelengthEncoder(wavelength_embed_dim)
+        self.dynamic_weights = DynamicPatchWeights(wavelength_embed_dim, embed_dim, decoder_dim,
+                                                   patch_size, hypernetwork_heads)
+        self.encoder_position = nn.Parameter(torch.zeros(1, self.num_patches, embed_dim))
+        encoder_layer = nn.TransformerEncoderLayer(embed_dim, encoder_heads, embed_dim * 4,
+                                                   batch_first=True, norm_first=True,
+                                                   dropout=0.0, activation="gelu")
+        self.encoder = nn.TransformerEncoder(encoder_layer, encoder_depth)
+        self.decoder_input = nn.Linear(embed_dim, decoder_dim)
+        self.decoder_position = nn.Parameter(torch.zeros(1, self.num_patches, decoder_dim))
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))
+        decoder_layer = nn.TransformerEncoderLayer(decoder_dim, decoder_heads, decoder_dim * 4,
+                                                   batch_first=True, norm_first=True,
+                                                   dropout=0.0, activation="gelu")
+        self.decoder = nn.TransformerEncoder(decoder_layer, decoder_depth)
+        nn.init.normal_(self.encoder_position, std=0.02)
+        nn.init.normal_(self.decoder_position, std=0.02)
+        nn.init.normal_(self.mask_token, std=0.02)
+
+    def patchify(self, images):
+        p = self.patch_size
+        batch, channels, height, width = images.shape
+        patches = images.reshape(batch, channels, height // p, p, width // p, p)
+        return patches.permute(0, 2, 4, 1, 3, 5).reshape(batch, self.num_patches,
+                                                        channels, p * p)
+
+    def unpatchify(self, patches):
+        p, side = self.patch_size, self.image_size // self.patch_size
+        batch, _, channels, _ = patches.shape
+        return patches.reshape(batch, side, side, channels, p, p).permute(
+            0, 3, 1, 4, 2, 5).reshape(batch, channels, self.image_size, self.image_size)
+
+    def random_mask(self, batch, ratio, device):
+        visible = max(1, int(self.num_patches * (1 - ratio)))
+        order = torch.rand(batch, self.num_patches, device=device).argsort(1)
+        visible_indices = order[:, :visible]
+        mask = torch.ones(batch, self.num_patches, dtype=torch.bool, device=device)
+        mask.scatter_(1, visible_indices, False)
+        return visible_indices, mask
+
+    def forward(self, images, wavelengths, mask_ratio=None):
+        if images.shape[-2:] != (self.image_size, self.image_size):
+            raise ValueError(f"Expected {self.image_size}x{self.image_size} images")
+        if images.shape[1] != wavelengths.numel():
+            raise ValueError("Image channels and wavelengths must have equal lengths")
+        wavelength_features = self.wavelength_encoder(wavelengths)
+        encoder_weight, encoder_bias, decoder_weight, decoder_bias = self.dynamic_weights(
+            wavelength_features)
+        tokens = F.conv2d(images, encoder_weight, encoder_bias,
+                          stride=self.patch_size).flatten(2).transpose(1, 2)
+        visible_indices, mask = self.random_mask(len(images), self.mask_ratio if mask_ratio is None
+                                                 else mask_ratio, images.device)
+        gather = visible_indices.unsqueeze(-1).expand(-1, -1, tokens.shape[-1])
+        visible_tokens = torch.gather(tokens + self.encoder_position, 1, gather)
+        encoded = self.encoder(visible_tokens)
+        decoded_tokens = self.mask_token.expand(len(images), self.num_patches, -1).clone()
+        visible_decoded = self.decoder_input(encoded).to(decoded_tokens.dtype)
+        decoded_tokens.scatter_(1, visible_indices.unsqueeze(-1).expand(-1, -1,
+                                decoded_tokens.shape[-1]), visible_decoded)
+        decoded = self.decoder(decoded_tokens + self.decoder_position)
+        predictions = torch.einsum("bnd,cpd->bncp", decoded, decoder_weight)
+        predictions = predictions + decoder_bias[None, None]
+        targets = self.patchify(images)
+        patch_error = (predictions - targets).square().mean((2, 3))
+        loss = (patch_error * mask).sum() / mask.sum().clamp_min(1)
+        return {"loss": loss, "reconstruction": self.unpatchify(predictions), "mask": mask,
+                "visible_indices": visible_indices, "features": encoded}
+
+
+__all__ = ["DOFA", "WavelengthEncoder", "DynamicPatchWeights"]

scripts/fake_data.py

@@ -0,0 +1,72 @@
+"""Generate deterministic multi-sensor data for DOFA pipeline validation."""
+
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def make_images(count, channels, size, seed):
+    rng = np.random.default_rng(seed)
+    y, x = np.mgrid[0:size, 0:size].astype(np.float32) / max(size - 1, 1)
+    images = np.empty((count, channels, size, size), dtype=np.float32)
+    for sample in range(count):
+        for channel in range(channels):
+            phase = rng.uniform(0, 2 * np.pi)
+            pattern = 0.45 + 0.25 * np.sin((channel + 1) * np.pi * x + phase)
+            pattern += 0.2 * np.cos((sample % 4 + 1) * np.pi * y - phase)
+            pattern += rng.normal(0, 0.02, (size, size))
+            images[sample, channel] = np.clip(pattern, 0, 1)
+    return images
+
+
+def main():
+    with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    data = config["data"]
+    root = ROOT / data["root"]
+    root.mkdir(parents=True, exist_ok=True)
+    for index, (name, modality) in enumerate(data["modalities"].items()):
+        if modality.get("wavelength_mode") == "synthetic_uniform":
+            wavelengths = modality["wavelength_start"] + np.arange(
+                modality["channels"], dtype=np.float32
+            ) * modality["wavelength_step"]
+        else:
+            wavelengths = np.asarray(modality["wavelengths"], dtype=np.float32)
+        if wavelengths.shape != (modality["channels"],):
+            raise ValueError(
+                f"{name} has {modality['channels']} channels but "
+                f"{wavelengths.size} wavelengths"
+            )
+        for split, count in (
+            ("train", data["train_samples_per_modality"]),
+            ("test", data["test_samples_per_modality"]),
+        ):
+            images = make_images(
+                count,
+                modality["channels"],
+                data["image_size"],
+                config["seed"] + index * 10 + (split == "test"),
+            )
+            output = root / f"{split}_{name}.npz"
+            np.savez_compressed(
+                output,
+                images=images,
+                wavelengths=wavelengths,
+                modality=np.asarray(name),
+                data_source=np.asarray("synthetic"),
+                protocol=np.asarray(data["protocol"]),
+                data_range=np.asarray(data["data_range"], dtype=np.float32),
+                wavelength_mode=np.asarray(modality.get("wavelength_mode", "configured")),
+            )
+            print(
+                f"generated={output.relative_to(ROOT)} shape={images.shape} "
+                f"modality={name} channels={modality['channels']} data_source=synthetic"
+            )
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,104 @@
+"""Reconstruct validated sensor NPZ files in bounded device batches."""
+
+import importlib.util
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_model_class():
+    spec = importlib.util.spec_from_file_location("dofa_model", ROOT / "model" / "dofa.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.DOFA
+
+
+def expected_wavelengths(modality):
+    if modality.get("wavelength_mode") == "synthetic_uniform":
+        return modality["wavelength_start"] + np.arange(modality["channels"], dtype=np.float32) * modality["wavelength_step"]
+    return np.asarray(modality["wavelengths"], dtype=np.float32)
+
+
+def scalar(archive, key, default=None):
+    if key not in archive:
+        if default is not None:
+            return default
+        raise ValueError(f"NPZ is missing required metadata: {key}")
+    if archive[key].ndim != 0:
+        raise ValueError(f"NPZ metadata {key} must be a scalar")
+    return archive[key].item()
+
+
+def validate_archive(archive, path, name, data_config):
+    if "images" not in archive or "wavelengths" not in archive:
+        raise ValueError(f"{path}: images and wavelengths are required")
+    images, wavelengths = archive["images"], archive["wavelengths"]
+    modality_config = data_config["modalities"][name]
+    if images.ndim != 4 or images.dtype != np.float32:
+        raise ValueError(f"{path}: images must be float32 NCHW")
+    expected_shape = (modality_config["channels"], data_config["image_size"], data_config["image_size"])
+    if images.shape[1:] != expected_shape:
+        raise ValueError(f"{path}: expected [N,{expected_shape[0]},224,224], got {images.shape}")
+    expected = expected_wavelengths(modality_config)
+    if wavelengths.shape != (modality_config["channels"],) or not np.issubdtype(wavelengths.dtype, np.floating) or not np.isfinite(wavelengths).all() or not np.allclose(wavelengths, expected, rtol=1e-5, atol=1e-6):
+        raise ValueError(f"{path}: wavelengths do not match configured sensor wavelengths")
+    modality = str(scalar(archive, "modality"))
+    protocol = str(scalar(archive, "protocol"))
+    source = str(scalar(archive, "data_source", "unknown"))
+    if modality != name or protocol != data_config["protocol"]:
+        raise ValueError(f"{path}: modality/protocol metadata does not match config")
+    data_range = float(scalar(archive, "data_range", data_config.get("data_range")))
+    if not np.isfinite(data_range) or data_range <= 0:
+        raise ValueError(f"{path}: PSNR requires a positive data_range metadata or config value")
+    return images, wavelengths.astype("float32", copy=False), protocol, source, data_range
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text(encoding="utf-8"))
+    checkpoint_path = ROOT / config["paths"]["checkpoint"]
+    if not checkpoint_path.exists():
+        raise FileNotFoundError("Missing checkpoint. Run `python scripts/train.py` first.")
+    device = torch.device("cuda" if torch.cuda.is_available() and config["runtime"]["device"] != "cpu" else "cpu")
+    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    if checkpoint.get("protocol") != config["data"]["protocol"]:
+        raise ValueError("Checkpoint and configured protocols do not match")
+    model = load_model_class()(**config["model"]).to(device)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    output_dir = ROOT / config["paths"]["inference_dir"]
+    output_dir.mkdir(parents=True, exist_ok=True)
+    data_root = ROOT / config["data"]["root"]
+    batch_size = config["runtime"]["inference_batch_size"]
+    torch.manual_seed(config["seed"])
+    for modality in config["data"]["modalities"]:
+        path = data_root / f"test_{modality}.npz"
+        if not path.exists():
+            raise FileNotFoundError(f"Missing test data: {path.relative_to(ROOT)}")
+        archive = np.load(path)
+        images, wavelengths, protocol, source, data_range = validate_archive(
+            archive, path, modality, config["data"])
+        reconstructions = np.empty_like(images)
+        masks = np.empty((len(images), model.num_patches), dtype=bool)
+        wavelength_tensor = torch.from_numpy(wavelengths).to(device)
+        with torch.inference_mode():
+            for start in range(0, len(images), batch_size):
+                stop = min(start + batch_size, len(images))
+                batch = torch.from_numpy(images[start:stop]).to(device)
+                output = model(batch, wavelength_tensor)
+                reconstructions[start:stop] = output["reconstruction"].cpu().numpy()
+                masks[start:stop] = output["mask"].cpu().numpy()
+                del batch, output
+        target = output_dir / f"{modality}_reconstruction.npz"
+        np.savez_compressed(target, inputs=images, reconstructions=reconstructions, masks=masks,
+                            wavelengths=wavelengths, modality=np.asarray(modality),
+                            data_source=np.asarray(source), protocol=np.asarray(protocol),
+                            data_range=np.asarray(data_range, dtype=np.float32))
+        print(f"output={target.relative_to(ROOT)} channels={images.shape[1]} batch_size={batch_size}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,83 @@
+"""Evaluate DOFA on masked pixels with explicit PSNR data ranges."""
+
+import json
+from pathlib import Path
+
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def scalar(archive, key):
+    if key not in archive or archive[key].ndim != 0:
+        raise ValueError(f"Inference metadata {key} must be present as a scalar")
+    return archive[key].item()
+
+
+def display(array):
+    selected = array[:3] if len(array) >= 3 else np.repeat(array[:1], 3, 0)
+    selected = selected.transpose(1, 2, 0)
+    return np.clip((selected - selected.min()) / max(np.ptp(selected), 1e-6), 0, 1)
+
+
+def pixel_mask(patch_mask, image_size, patch_size):
+    side = image_size // patch_size
+    return np.repeat(np.repeat(patch_mask.reshape(side, side), patch_size, 0), patch_size, 1)
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text(encoding="utf-8"))
+    input_dir = ROOT / config["paths"]["inference_dir"]
+    output_dir = ROOT / config["paths"]["evaluation_dir"]
+    output_dir.mkdir(parents=True, exist_ok=True)
+    metrics, sources, protocols = {}, set(), set()
+    for modality, modality_config in config["data"]["modalities"].items():
+        path = input_dir / f"{modality}_reconstruction.npz"
+        if not path.is_file():
+            raise FileNotFoundError("Run inference before result evaluation")
+        archive = np.load(path)
+        inputs, reconstructions, masks = archive["inputs"], archive["reconstructions"], archive["masks"]
+        protocol = str(scalar(archive, "protocol"))
+        source = str(scalar(archive, "data_source"))
+        output_modality = str(scalar(archive, "modality"))
+        data_range = float(scalar(archive, "data_range"))
+        expected_shape = (modality_config["channels"], config["data"]["image_size"], config["data"]["image_size"])
+        if inputs.ndim != 4 or inputs.shape[1:] != expected_shape or reconstructions.shape != inputs.shape:
+            raise ValueError(f"{path}: invalid reconstruction NCHW shape")
+        if protocol != config["data"]["protocol"] or output_modality != modality:
+            raise ValueError(f"{path}: protocol/modality metadata does not match config")
+        if not np.isfinite(data_range) or data_range <= 0:
+            raise ValueError(f"{path}: PSNR requires a positive data_range")
+        sources.add(source)
+        protocols.add(protocol)
+        expanded = np.stack([pixel_mask(mask, config["data"]["image_size"],
+                            config["model"]["patch_size"]) for mask in masks])[:, None]
+        errors = (inputs - reconstructions)[np.broadcast_to(expanded, inputs.shape)]
+        mse, mae = float(np.mean(errors**2)), float(np.mean(np.abs(errors)))
+        metrics[modality] = {"masked_mse": mse, "masked_mae": mae,
+                             "masked_psnr_db": float(10 * np.log10(data_range**2 / max(mse, 1e-12))),
+                             "psnr_data_range": data_range, "channels": int(inputs.shape[1]),
+                             "masked_fraction": float(expanded.mean()), "data_source": source,
+                             "protocol": protocol}
+        masked_input = inputs[0].copy()
+        masked_input[:, expanded[0, 0].astype(bool)] = 0
+        figure, axes = plt.subplots(1, 3, figsize=(10, 3))
+        for axis, image, title in zip(axes, (inputs[0], masked_input, reconstructions[0]),
+                                      ("target", "masked input", "reconstruction")):
+            axis.imshow(display(image)); axis.set_title(f"{modality}: {title}"); axis.axis("off")
+        figure.tight_layout()
+        figure.savefig(output_dir / f"{modality}_comparison.png", dpi=120)
+        plt.close(figure)
+    payload = {"modalities": metrics, "data_sources": sorted(sources),
+               "protocols": sorted(protocols), "protocol": config["data"]["protocol"],
+               "metric_scope": "masked_pixels_only"}
+    (output_dir / "metrics.json").write_text(json.dumps(payload, indent=2) + "\n")
+    print(json.dumps(payload, indent=2)); print(f"evaluation={output_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,167 @@
+"""Train compact DOFA across strictly validated, variable-channel sensors."""
+
+import importlib.util
+import json
+import os
+import random
+from contextlib import nullcontext
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch import distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_model_class():
+    spec = importlib.util.spec_from_file_location("dofa_model", ROOT / "model" / "dofa.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.DOFA
+
+
+def expected_wavelengths(modality):
+    if modality.get("wavelength_mode") == "synthetic_uniform":
+        return modality["wavelength_start"] + np.arange(modality["channels"], dtype=np.float32) * modality["wavelength_step"]
+    return np.asarray(modality["wavelengths"], dtype=np.float32)
+
+
+def scalar(archive, key, default=None):
+    if key not in archive:
+        if default is not None:
+            return default
+        raise ValueError(f"NPZ is missing required metadata: {key}")
+    if archive[key].ndim != 0:
+        raise ValueError(f"NPZ metadata {key} must be a scalar")
+    return archive[key].item()
+
+
+class SensorDataset(Dataset):
+    def __init__(self, path, name, data_config):
+        archive = np.load(path)
+        if "images" not in archive or "wavelengths" not in archive:
+            raise ValueError(f"{path} must contain images and wavelengths")
+        self.images = archive["images"]
+        wavelengths = archive["wavelengths"]
+        modality = data_config["modalities"][name]
+        expected = expected_wavelengths(modality)
+        if self.images.ndim != 4:
+            raise ValueError(f"{path}: images must be NCHW")
+        if self.images.dtype != np.float32:
+            raise ValueError(f"{path}: images must use float32")
+        if self.images.shape[1:] != (modality["channels"], data_config["image_size"], data_config["image_size"]):
+            raise ValueError(f"{path}: image shape does not match configured channels/224x224")
+        if wavelengths.shape != (modality["channels"],) or not np.issubdtype(wavelengths.dtype, np.floating):
+            raise ValueError(f"{path}: wavelengths must be a floating [C] array")
+        if not np.isfinite(wavelengths).all() or not np.allclose(wavelengths, expected, rtol=1e-5, atol=1e-6):
+            raise ValueError(f"{path}: wavelengths do not match configured sensor wavelengths")
+        self.modality = str(scalar(archive, "modality"))
+        self.protocol = str(scalar(archive, "protocol"))
+        self.data_source = str(scalar(archive, "data_source", "unknown"))
+        if self.modality != name:
+            raise ValueError(f"{path}: modality {self.modality} does not match {name}")
+        if self.protocol != data_config["protocol"]:
+            raise ValueError(f"{path}: protocol {self.protocol} does not match config")
+        self.wavelengths = torch.from_numpy(wavelengths.astype("float32", copy=False))
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        return torch.from_numpy(self.images[index])
+
+
+def reduced_average(total, count, device, distributed):
+    values = torch.tensor([total, count], dtype=torch.float64, device=device)
+    if distributed:
+        dist.all_reduce(values, op=dist.ReduceOp.SUM)
+    if values[1].item() == 0:
+        raise RuntimeError("Training processed no batches")
+    return (values[0] / values[1]).item()
+
+
+def main():
+    config = yaml.safe_load((ROOT / "conf/config.yaml").read_text(encoding="utf-8"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    global_rank = int(os.environ.get("RANK", "0"))
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    distributed = world_size > 1
+    if distributed:
+        dist.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
+    requested = config["runtime"]["device"]
+    use_accelerator = torch.cuda.is_available() and requested != "cpu"
+    device = torch.device(f"cuda:{local_rank}" if use_accelerator else "cpu")
+    if use_accelerator:
+        torch.cuda.set_device(local_rank)
+    seed = config["seed"] + global_rank
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+    data_root = ROOT / config["data"]["root"]
+    loaders = []
+    for modality in config["data"]["modalities"]:
+        path = data_root / f"train_{modality}.npz"
+        if not path.exists():
+            raise FileNotFoundError(f"Missing training data: {path.relative_to(ROOT)}. Run `python scripts/fake_data.py` first.")
+        dataset = SensorDataset(path, modality, config["data"])
+        sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+        loader = DataLoader(dataset, batch_size=config["training"]["batch_size"],
+                            shuffle=sampler is None, sampler=sampler,
+                            num_workers=config["training"]["num_workers"])
+        loaders.append((dataset, loader, sampler))
+
+    model = load_model_class()(**config["model"]).to(device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if use_accelerator else None)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=config["training"]["learning_rate"],
+                                  weight_decay=config["training"]["weight_decay"])
+    amp = bool(config["training"]["amp"] and use_accelerator)
+    scaler = torch.amp.GradScaler("cuda", enabled=amp)
+    autocast = (lambda: torch.amp.autocast("cuda", enabled=True)) if amp else nullcontext
+    history = []
+    for epoch in range(config["training"]["epochs"]):
+        model.train()
+        records = {}
+        for dataset, loader, sampler in loaders:
+            if sampler is not None:
+                sampler.set_epoch(epoch)
+            total = count = 0
+            for images in loader:
+                optimizer.zero_grad(set_to_none=True)
+                with autocast():
+                    output = model(images.to(device), dataset.wavelengths.to(device))
+                scaler.scale(output["loss"]).backward()
+                scaler.step(optimizer)
+                scaler.update()
+                total += output["loss"].item()
+                count += 1
+            records[dataset.modality] = reduced_average(total, count, device, distributed)
+        history.append({"epoch": epoch + 1, "reconstruction_loss": records})
+        if global_rank == 0:
+            print(f"epoch={epoch + 1} " + " ".join(f"{name}={value:.6f}" for name, value in records.items()))
+
+    if global_rank == 0:
+        checkpoint = ROOT / config["paths"]["checkpoint"]
+        metrics = ROOT / config["paths"]["training_metrics"]
+        checkpoint.parent.mkdir(parents=True, exist_ok=True)
+        metrics.parent.mkdir(parents=True, exist_ok=True)
+        base_model = model.module if hasattr(model, "module") else model
+        sources = sorted({dataset.data_source for dataset, _, _ in loaders})
+        torch.save({"model": base_model.state_dict(), "config": config,
+                    "protocol": config["data"]["protocol"], "data_sources": sources}, checkpoint)
+        metrics.write_text(json.dumps({"history": history,
+                           "modalities": list(config["data"]["modalities"]),
+                           "protocol": config["data"]["protocol"],
+                           "data_sources": sources}, indent=2) + "\n", encoding="utf-8")
+        print(f"checkpoint={checkpoint.relative_to(ROOT)}")
+    if distributed:
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()