Add Chinese Hugging Face model card and configuration

.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

README.md

@@ -0,0 +1,196 @@
+---
+frameworks: PyTorch
+language:
+- zh
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Remote Sensing
+- Masked Image Modeling
+- Sentinel-2
+- SpectralGPT
+- arxiv:2311.07113
+tasks: []
+datasets: []
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">SpectralGPT</span>
+  </strong>
+</p>
+
+# 模型介绍
+
+SpectralGPT 通过三维空间-光谱掩码建模和多目标重建，从大规模多光谱遥感影像中学习通用表征，主要用于少标签条件下的场景分类、语义分割和变化检测。
+
+论文：SpectralGPT: Spectral Remote Sensing Foundation Model
+
+https://arxiv.org/abs/2311.07113
+
+# 模型描述
+
+SpectralGPT 由中国科学院空天信息创新研究院等机构提出，使用 fMoW-S2 和 BigEarthNet-S2 的 12 波段 Sentinel-2 影像训练。模型适用于单标签或多标签场景分类、语义分割和变化检测等遥感任务。
+
+# 适用场景
+
+| 场景 | 说明 |
+| :---: | :--- |
+| 多光谱掩码预训练 | 使用归一化的 12 波段 Sentinel-2 多光谱影像训练 SpectralGPT。 |
+| 遥感场景理解 | 使用学习到的多光谱表征进行场景分类、语义分割和变化检测。 |
+| 本地快速验证 | 使用虚拟数据检查数据读取、掩码重建训练、推理、评估和可视化流程。 |
+| Hugging Face / OneCode 运行 | 作为独立模型包下载后直接安装依赖并运行脚本。 |
+| 多卡训练 | 通过 `torchrun` 启动多进程训练。 |
+
+# 使用说明
+
+## 1. OneCode 使用
+
+可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：
+
+[点击体验智能化一键式 AI4S 编程](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. 手动安装使用
+
+**硬件要求**
+
+- 推荐使用 GPU 或 DCU 运行。
+- CPU 可以用于导入和小配置连通性验证，完整训练和推理速度较慢。
+- DCU 用户需要预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。
+
+### 下载模型包
+
+```bash
+hf download OneScience-Group/SpectralGPT --local-dir ./SpectralGPT
+cd SpectralGPT
+```
+
+### 安装运行环境
+
+**DCU环境**
+
+```bash
+# 请首先激活 DTK 及 CONDA
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# 支持 uv 安装
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU环境**
+
+```bash
+# 请首先激活 CONDA
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# 支持 uv 安装
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### 训练数据介绍
+
+论文使用 Sentinel-2 的 12 个主波段并排除 B10，各波段缩放到 `[0,1]`。fMoW-S2 包含 882,779 张影像，其中 712,874 张用于第一阶段预训练；BigEarthNet-S2 包含 590,326 张影像，其中 354,196 张用于后续预训练。OneScience 社区目前未提供可直接用于本仓库训练的 fMoW-S2、BigEarthNet-S2 或 EuroSAT 数据，因此本仓库默认使用 `scripts/fake_data.py` 生成 8 张小规模虚拟多光谱影像进行流程验证。
+
+EuroSAT 官方数据资源：
+
+```text
+官方 Zenodo 数据页：
+https://zenodo.org/records/7711810
+
+多光谱数据下载：
+https://madm.dfki.de/files/sentinel/EuroSATallBands.zip
+```
+
+解压 EuroSAT 多光谱数据后，可使用以下命令排除 B10、缩放到 `[0,1]` 并转换为本仓库使用的 12 波段 NPZ 数据：
+
+```bash
+python scripts/fake_data.py --real-dir ./data/EuroSAT_MS
+```
+
+虚拟数据保存为 `data/fake_spectralgpt.npz`，只用于验证数据读取、掩码重建训练、检查点保存与加载、推理、评估和可视化流程，不代表真实 Sentinel-2 影像，也不能用于复现论文指标。真实数据训练不应先运行不带 `--real-dir` 的虚拟数据命令，并需要使用一致的波段顺序、空间尺寸和数值缩放方式。
+
+### 训练
+
+使用虚拟数据快速验证：
+
+```bash
+python scripts/fake_data.py
+python scripts/train.py
+```
+
+使用真实 EuroSAT 多光谱 TIFF 训练：
+
+```bash
+python scripts/fake_data.py --real-dir ./data/EuroSAT_MS
+python scripts/train.py
+```
+
+虚拟数据和真实数据共用同一个训练脚本。需要调整训练规模时，可修改 `conf/config.yaml` 中的 `data.samples`、`model.image_size`、`training.epochs` 和 `training.batch_size`。
+
+单卡：
+
+```bash
+python scripts/train.py
+```
+
+多卡：
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+训练输出：
+
+```text
+result/checkpoints/best.pth
+result/checkpoints/last.pth
+每轮输出 reconstruction_loss。
+```
+
+训练日志和检查点中的 `data_source` 与 `protocol` 用于标识当前使用的是虚拟数据还是真实数据���使用虚拟数据生成的检查点只用于流程验证，不具备真实遥感表征能力。
+
+### 训练权重
+
+本仓库将在 `weight/` 文件夹内提供基于 Sentinel-2 多光谱遥感数据训练的权重，权重文件即将上传，预计将于近期完成。
+
+### 推理
+
+```bash
+python scripts/inference.py
+```
+
+推理读取训练数据路径配置和训练生成的检查点，并输出多光谱影像重建结果。
+
+推理结果输出到：
+
+```text
+result/output/reconstruction.npz
+```
+
+### 评估和可视化
+
+```bash
+python scripts/result.py
+```
+
+评估和可视化输出到：
+
+```text
+result/output/metrics.json
+result/output/reconstruction.ppm
+```
+
+指标包括重建 MSE、MAE、PSNR 和逐波段 RMSE；可视化并排显示输入与重建的假彩色合成图。使用虚拟数据时，评估结果只用于确认代码流程能够运行，不展示或报告虚拟数据数值；使用真实数据时，指标表示当前真实影像上的掩码重建结果。
+
+# OneScience 官方信息
+
+| 平台 | OneScience 主仓库 | Skills 仓库 |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# 引用与许可证
+
+- 本仓库为 SpectralGPT 原始论文的复现版本。
+- 本模型包采用 Apache License 2.0；原始数据与论文作者发布的模型权重遵循各自来源的许可条款。

conf/config.yaml

@@ -0,0 +1,28 @@
+model:
+  image_size: 24
+  in_channels: 12
+  patch_size: 8
+  spectral_patch_size: 3
+  embed_dim: 48
+  encoder_depth: 2
+  encoder_heads: 4
+  decoder_dim: 32
+  decoder_depth: 1
+  decoder_heads: 4
+  mask_ratio: 0.90
+  spectral_loss_weight: 1.0
+data:
+  path: ./data/fake_spectralgpt.npz
+  samples: 8
+  protocol: synthetic_sentinel2_npz
+training:
+  epochs: 2
+  batch_size: 2
+  learning_rate: 0.0001
+  weight_decay: 0.05
+  save_dir: ./result/checkpoints
+  checkpoint: ./result/checkpoints/best.pth
+runtime:
+  device: auto
+  seed: 42
+  output_dir: ./result/output

config.json

@@ -0,0 +1,47 @@
+{
+  "model_name": "SpectralGPT",
+  "model_type": "spectralgpt",
+  "architectures": [
+    "SpectralGPT"
+  ],
+  "framework": "PyTorch",
+  "domain": "earth-science",
+  "task": "remote-sensing-masked-image-modeling",
+  "implementation": {
+    "entry_point": "model/spectralgpt.py",
+    "scope": "compact 12-band Sentinel-2 masked autoencoder reproduction"
+  },
+  "architecture": {
+    "family": "spectral-spatial masked autoencoder",
+    "input_format": "NCHW multispectral images",
+    "input_channels": 12,
+    "image_size": 24,
+    "patch_size": 8,
+    "spectral_patch_size": 3,
+    "embed_dim": 48,
+    "encoder_depth": 2,
+    "decoder_dim": 32,
+    "decoder_depth": 1,
+    "mask_ratio": 0.9,
+    "training_objective": "masked reconstruction with spectral loss"
+  },
+  "data": {
+    "datasets": [
+      "fMoW-S2",
+      "BigEarthNet-S2"
+    ],
+    "official_paper": "https://arxiv.org/abs/2311.07113",
+    "protocol": "synthetic_sentinel2_npz",
+    "synthetic_samples": 8
+  },
+  "metrics": [
+    "mse",
+    "mae",
+    "psnr_db",
+    "per_band_rmse"
+  ],
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/spectralgpt.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "remote_sensing_masked_image_modeling",
+  "model": "SpectralGPT",
+  "input_format": "NCHW_NPZ",
+  "protocol": "synthetic_sentinel2_npz",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/spectralgpt.py

@@ -0,0 +1,112 @@
+import torch
+from torch import nn
+
+
+class SpectralGPT(nn.Module):
+    """Compact SpectralGPT masked autoencoder for 12-band spectral images."""
+
+    def __init__(self, image_size=24, in_channels=12, patch_size=8,
+                 spectral_patch_size=3, embed_dim=48, encoder_depth=2,
+                 encoder_heads=4, decoder_dim=32, decoder_depth=1,
+                 decoder_heads=4, mask_ratio=0.9, spectral_loss_weight=1.0):
+        super().__init__()
+        if image_size % patch_size or in_channels % spectral_patch_size:
+            raise ValueError("Image and spectral dimensions must be divisible by token sizes")
+        self.image_size = image_size
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.spectral_patch_size = spectral_patch_size
+        self.spatial_tokens = (image_size // patch_size) ** 2
+        self.spectral_tokens = in_channels // spectral_patch_size
+        self.num_tokens = self.spatial_tokens * self.spectral_tokens
+        self.token_pixels = patch_size * patch_size * spectral_patch_size
+        self.mask_ratio = mask_ratio
+        self.spectral_loss_weight = spectral_loss_weight
+
+        self.patch_embed = nn.Conv3d(
+            1, embed_dim,
+            kernel_size=(spectral_patch_size, patch_size, patch_size),
+            stride=(spectral_patch_size, patch_size, patch_size),
+        )
+        self.spatial_pos = nn.Parameter(torch.zeros(1, self.spatial_tokens, embed_dim))
+        self.spectral_pos = nn.Parameter(torch.zeros(1, self.spectral_tokens, embed_dim))
+        encoder_layer = nn.TransformerEncoderLayer(
+            embed_dim, encoder_heads, embed_dim * 4, batch_first=True, norm_first=True
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, encoder_depth)
+        self.encoder_norm = nn.LayerNorm(embed_dim)
+        self.decoder_embed = nn.Linear(embed_dim, decoder_dim)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))
+        self.decoder_pos = nn.Linear(embed_dim, decoder_dim, bias=False)
+        decoder_layer = nn.TransformerEncoderLayer(
+            decoder_dim, decoder_heads, decoder_dim * 4, batch_first=True, norm_first=True
+        )
+        self.decoder = nn.TransformerEncoder(decoder_layer, decoder_depth)
+        self.decoder_norm = nn.LayerNorm(decoder_dim)
+        self.decoder_pred = nn.Linear(decoder_dim, self.token_pixels)
+        nn.init.normal_(self.spatial_pos, std=0.02)
+        nn.init.normal_(self.spectral_pos, std=0.02)
+        nn.init.normal_(self.mask_token, std=0.02)
+
+    def _positions(self):
+        return (self.spatial_pos[:, None] + self.spectral_pos[:, :, None]).reshape(
+            1, self.num_tokens, -1
+        )
+
+    def patchify(self, images):
+        p, k = self.patch_size, self.spectral_patch_size
+        n, c, h, w = images.shape
+        if (c, h, w) != (self.in_channels, self.image_size, self.image_size):
+            raise ValueError(f"Expected [N,{self.in_channels},{self.image_size},{self.image_size}]")
+        x = images.reshape(n, c // k, k, h // p, p, w // p, p)
+        x = x.permute(0, 1, 3, 5, 2, 4, 6)
+        return x.reshape(n, self.num_tokens, self.token_pixels)
+
+    def unpatchify(self, tokens):
+        p, k = self.patch_size, self.spectral_patch_size
+        n = tokens.shape[0]
+        s = self.image_size // p
+        x = tokens.reshape(n, self.spectral_tokens, s, s, k, p, p)
+        x = x.permute(0, 1, 4, 2, 5, 3, 6)
+        return x.reshape(n, self.in_channels, self.image_size, self.image_size)
+
+    @staticmethod
+    def random_masking(tokens, mask_ratio):
+        n, length, dim = tokens.shape
+        keep = max(1, int(length * (1.0 - mask_ratio)))
+        order = torch.argsort(torch.rand(n, length, device=tokens.device), dim=1)
+        restore = torch.argsort(order, dim=1)
+        keep_ids = order[:, :keep]
+        visible = torch.gather(tokens, 1, keep_ids.unsqueeze(-1).expand(-1, -1, dim))
+        mask = torch.ones(n, length, device=tokens.device)
+        mask[:, :keep] = 0
+        mask = torch.gather(mask, 1, restore)
+        return visible, mask, restore
+
+    def forward(self, images, mask_ratio=None):
+        ratio = self.mask_ratio if mask_ratio is None else mask_ratio
+        embedded = self.patch_embed(images.unsqueeze(1)).flatten(2).transpose(1, 2)
+        positions = self._positions()
+        visible, mask, restore = self.random_masking(embedded + positions, ratio)
+        latent = self.encoder_norm(self.encoder(visible))
+        decoded_visible = self.decoder_embed(latent)
+        missing = self.num_tokens - decoded_visible.shape[1]
+        full = torch.cat([decoded_visible, self.mask_token.expand(images.shape[0], missing, -1)], 1)
+        full = torch.gather(full, 1, restore.unsqueeze(-1).expand(-1, -1, full.shape[-1]))
+        prediction = self.decoder_pred(self.decoder_norm(self.decoder(full + self.decoder_pos(positions))))
+        target = self.patchify(images)
+        token_error = (prediction - target).pow(2).mean(-1)
+        token_loss = (token_error * mask).sum() / mask.sum().clamp_min(1)
+        n = images.shape[0]
+        predicted_grid = prediction.reshape(n, self.spectral_tokens, self.spatial_tokens, -1)
+        target_grid = target.reshape(n, self.spectral_tokens, self.spatial_tokens, -1)
+        spectral_loss = (predicted_grid[:, 1:] - target_grid[:, 1:]).pow(2).mean()
+        loss = token_loss + self.spectral_loss_weight * spectral_loss
+        return {
+            "loss": loss,
+            "token_loss": token_loss,
+            "spectral_loss": spectral_loss,
+            "prediction": prediction,
+            "mask": mask,
+            "reconstruction": self.unpatchify(prediction),
+        }

scripts/fake_data.py

@@ -0,0 +1,67 @@
+import argparse
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+def real_images(directory, size, samples):
+    import tifffile
+
+    paths = sorted(Path(directory).rglob("*.tif"))[:samples]
+    if not paths:
+        raise FileNotFoundError(f"No TIFF files found under {directory}")
+    images = []
+    for path in paths:
+        image = tifffile.imread(path).astype(np.float32)
+        if image.ndim != 3:
+            raise ValueError(f"Expected a 13-band TIFF, got {image.shape} from {path}")
+        if image.shape[0] == 13:
+            image = image.transpose(1, 2, 0)
+        if image.shape[-1] != 13:
+            raise ValueError(f"Expected 13 Sentinel-2 bands, got {image.shape} from {path}")
+        image = np.delete(image, 10, axis=-1)  # Exclude Sentinel-2 B10 as in the paper.
+        y = np.linspace(0, image.shape[0] - 1, size).round().astype(int)
+        x = np.linspace(0, image.shape[1] - 1, size).round().astype(int)
+        image = image[y][:, x]
+        images.append(np.clip(image / 10000.0, 0, 1).transpose(2, 0, 1))
+    return np.asarray(images, dtype=np.float32)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate compact 12-band spectral data")
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--real-dir", help="Convert official 13-band Sentinel-2 TIFF files")
+    args = parser.parse_args()
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    size = config["model"]["image_size"]
+    samples = config["data"]["samples"]
+    if args.real_dir:
+        images = real_images(args.real_dir, size, samples)
+        source = "real"
+        protocol = "real_sentinel2_npz"
+    else:
+        rng = np.random.default_rng(config["runtime"]["seed"])
+        y, x = np.mgrid[0:size, 0:size].astype(np.float32) / max(size - 1, 1)
+        images = []
+        for index in range(samples):
+            phase = rng.uniform(0, 2 * np.pi)
+            bands = []
+            for band in range(12):
+                pattern = 0.45 + 0.22 * np.sin((band + 1) * x + phase)
+                pattern += 0.18 * np.cos((band / 3 + 1) * y - phase)
+                pattern += rng.normal(0, 0.025, (size, size))
+                bands.append(np.clip(pattern, 0, 1))
+            images.append(bands)
+        images = np.asarray(images, dtype=np.float32)
+        source = "synthetic"
+        protocol = config["data"]["protocol"]
+    output = Path(config["data"]["path"])
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, images=images, data_source=np.asarray(source), protocol=np.asarray(protocol))
+    print(f"saved: {output} shape={images.shape} data_source={source} protocol={protocol}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,58 @@
+import argparse
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+import yaml
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.spectralgpt import SpectralGPT
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run SpectralGPT reconstruction")
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--checkpoint")
+    args = parser.parse_args()
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    requested = config["runtime"]["device"]
+    device = torch.device("cuda" if torch.cuda.is_available() and requested != "cpu" else "cpu")
+    torch.manual_seed(config["runtime"]["seed"])
+    model = SpectralGPT(**config["model"]).to(device)
+    checkpoint_path = args.checkpoint or config["training"]["checkpoint"]
+    if not Path(checkpoint_path).exists():
+        raise FileNotFoundError(
+            f"Missing checkpoint: {checkpoint_path}. Run `python scripts/train.py` first."
+        )
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    model.load_state_dict(checkpoint["model"])
+    model.eval()
+    data_path = Path(config["data"]["path"])
+    if not data_path.exists():
+        raise FileNotFoundError(
+            f"Missing inference data: {data_path}. Run `python scripts/fake_data.py` first."
+        )
+    with np.load(data_path) as data:
+        images = torch.from_numpy(data["images"]).to(device)
+        data_source = str(data["data_source"]) if "data_source" in data.files else "unknown"
+        protocol = str(data["protocol"]) if "protocol" in data.files else "unknown"
+    with torch.no_grad():
+        output = model(images)
+    output_dir = Path(config["runtime"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output_dir / "reconstruction.npz",
+                        inputs=images.cpu().numpy(),
+                        reconstructions=output["reconstruction"].cpu().numpy(),
+                        masks=output["mask"].cpu().numpy(),
+                        data_source=np.asarray(data_source),
+                        protocol=np.asarray(protocol))
+    print(
+        f"saved: {output_dir / 'reconstruction.npz'} "
+        f"data_source={data_source} protocol={protocol}"
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,51 @@
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+def rgb(image):
+    array = np.clip(image[[3, 2, 1]], 0, 1).transpose(1, 2, 0)
+    return (array * 255).astype(np.uint8)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate and visualize reconstruction")
+    parser.add_argument("--config", default="conf/config.yaml")
+    args = parser.parse_args()
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    output_dir = Path(config["runtime"]["output_dir"])
+    reconstruction_path = output_dir / "reconstruction.npz"
+    if not reconstruction_path.exists():
+        raise FileNotFoundError(
+            f"Missing inference output: {reconstruction_path}. "
+            "Run `python scripts/inference.py` first."
+        )
+    with np.load(reconstruction_path) as data:
+        inputs = data["inputs"]
+        reconstructions = data["reconstructions"]
+        data_source = str(data["data_source"]) if "data_source" in data.files else "unknown"
+        protocol = str(data["protocol"]) if "protocol" in data.files else "unknown"
+    mse = float(np.mean((inputs - reconstructions) ** 2))
+    mae = float(np.mean(np.abs(inputs - reconstructions)))
+    psnr = float(-10 * np.log10(max(mse, 1e-12)))
+    spectral_rmse = np.sqrt(np.mean((inputs - reconstructions) ** 2, axis=(0, 2, 3)))
+    metrics = {"mse": mse, "mae": mae, "psnr_db": psnr,
+               "data_source": data_source, "protocol": protocol,
+               "per_band_rmse": spectral_rmse.tolist()}
+    with open(output_dir / "metrics.json", "w", encoding="utf-8") as handle:
+        json.dump(metrics, handle, indent=2)
+    panel = np.concatenate([rgb(inputs[0]), rgb(reconstructions[0])], axis=1)
+    with open(output_dir / "reconstruction.ppm", "wb") as handle:
+        handle.write(f"P6\n{panel.shape[1]} {panel.shape[0]}\n255\n".encode("ascii"))
+        handle.write(panel.tobytes())
+    print(json.dumps(metrics, indent=2))
+    print(f"saved: {output_dir / 'metrics.json'}")
+    print(f"saved: {output_dir / 'reconstruction.ppm'}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,116 @@
+import argparse
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+import torch
+import yaml
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+from model.spectralgpt import SpectralGPT
+
+
+class SpectralDataset(Dataset):
+    def __init__(self, path):
+        with np.load(path) as data:
+            if "images" not in data.files:
+                raise ValueError(f"Dataset {path} is missing the images array")
+            self.images = data["images"].copy()
+            self.data_source = str(data["data_source"]) if "data_source" in data.files else "unknown"
+            self.protocol = str(data["protocol"]) if "protocol" in data.files else "unknown"
+
+    def __len__(self):
+        return len(self.images)
+
+    def __getitem__(self, index):
+        return torch.from_numpy(self.images[index])
+
+
+def load_config(path):
+    with open(path, encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def build_model(config):
+    return SpectralGPT(**config["model"])
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train compact SpectralGPT")
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--data")
+    parser.add_argument("--epochs", type=int)
+    args = parser.parse_args()
+    config = load_config(args.config)
+    distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    if distributed:
+        torch.distributed.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")
+    torch.manual_seed(config["runtime"]["seed"] + local_rank)
+    data_path = Path(args.data or config["data"]["path"])
+    if not data_path.exists():
+        raise FileNotFoundError(
+            f"Missing training data: {data_path}. "
+            "Run `python scripts/fake_data.py` for a synthetic connectivity test."
+        )
+    dataset = SpectralDataset(data_path)
+    if dataset.images.ndim != 4 or tuple(dataset.images.shape[1:]) != (
+        config["model"]["in_channels"], config["model"]["image_size"], config["model"]["image_size"]
+    ):
+        raise ValueError(
+            f"Expected images shaped [N,{config['model']['in_channels']},"
+            f"{config['model']['image_size']},{config['model']['image_size']}], "
+            f"got {dataset.images.shape}"
+        )
+    if local_rank == 0:
+        print(
+            f"data_source={dataset.data_source} protocol={dataset.protocol} "
+            f"samples={len(dataset)}"
+        )
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"],
+                        sampler=sampler, shuffle=sampler is None)
+    model = build_model(config).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"], betas=(0.9, 0.95))
+    save_dir = Path(config["training"]["save_dir"])
+    best = float("inf")
+    for epoch in range(args.epochs or config["training"]["epochs"]):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train()
+        losses = []
+        for images in loader:
+            output = model(images.to(device))
+            optimizer.zero_grad()
+            output["loss"].backward()
+            optimizer.step()
+            losses.append(output["loss"].item())
+        mean_loss = float(np.mean(losses))
+        if local_rank == 0:
+            print(f"epoch={epoch + 1} reconstruction_loss={mean_loss:.6f}")
+            save_dir.mkdir(parents=True, exist_ok=True)
+            state = model.module.state_dict() if distributed else model.state_dict()
+            checkpoint = {"model": state, "epoch": epoch + 1, "loss": mean_loss,
+                          "data_source": dataset.data_source, "protocol": dataset.protocol,
+                          "format": "spectralgpt-compact-v1"}
+            torch.save(checkpoint, save_dir / "last.pth")
+            if mean_loss < best:
+                best = mean_loss
+                torch.save(checkpoint, save_dir / "best.pth")
+    if local_rank == 0:
+        print(f"saved: {save_dir / 'best.pth'}")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()