Publish SatlasPretrain reproduction

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin 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
-*.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
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel 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
+*.wasm 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,181 @@
+---
+frameworks: PyTorch
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Remote Sensing
+- Multi-Task Image Understanding
+- SatlasPretrain
+- arxiv:2211.15660
+tasks: []
+datasets:
+- allenai/satlas-pretrain
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">SatlasPretrain</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+SatlasPretrain is a large-scale remote sensing pre-training approach for learning from multi-temporal imagery and seven label types. Its unified SatlasNet model supports semantic segmentation, regression, point detection, polygon instances, polylines, object properties, and image classification for Earth observation and downstream remote sensing applications.
+
+Paper: SatlasPretrain: A Large-Scale Dataset for Remote Sensing Image Understanding
+
+https://arxiv.org/abs/2211.15660
+
+# Model Description
+
+SatlasPretrain was proposed by the Allen Institute for AI and trained on Sentinel-2 and NAIP imagery with 302 million labels in 137 categories and seven label types.
+The model is suitable for multi-temporal remote sensing segmentation, regression, detection, property prediction, and classification tasks.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Remote sensing multi-task training | Train SatlasPretrain with multi-temporal imagery and segmentation, regression, detection, property, and classification labels. |
+| Local quick validation | Use synthetic data to validate loading, multi-task training, inference, evaluation, and qualitative visualization. |
+| Hugging Face / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
+| Multi-GPU training | Launch distributed multi-process training with `torchrun`. |
+
+# Usage Guide
+
+## 1. OneCode Usage
+
+Experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- CPU can be used for import and small-scale connectivity verification; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching the cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/SatlasPretrain --local-dir ./SatlasPretrain
+cd SatlasPretrain
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Please activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation is supported
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Please activate CONDA first
+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 installation is supported
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+### Training Data Introduction
+
+The original SatlasPretrain dataset is published by the Allen Institute for AI and organized as Web-Mercator tiles. Sentinel-2, Sentinel-1, and NAIP imagery and bands are stored as PNG files. Point, polygon, polyline, property, and classification labels are stored in `vector.json`; semantic segmentation and regression labels are grayscale PNG files; split and image-time metadata are JSON files.
+
+Official data resources:
+
+- https://github.com/allenai/satlas/blob/main/SatlasPretrain.md
+- https://huggingface.co/datasets/allenai/satlas-pretrain
+
+The original dataset is not included in this model repository. The provided scripts use synthetic NPZ data to validate the full training, inference, evaluation, and visualization pipeline. This NPZ protocol is only a compact pipeline check and is not the original SatlasPretrain format.
+
+Generate synthetic data:
+
+```bash
+python scripts/fake_data.py
+```
+
+The generated `data/fake_satlaspretrain.npz` contains four samples with `images` in `[B, T, C, H, W]` layout and labels for all seven task heads. Verify that `conf/config.yaml` points to this file before training.
+
+### Training
+
+Single GPU or CPU:
+
+```bash
+python scripts/train.py
+```
+
+Multi-GPU:
+
+```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
+```
+
+Training outputs:
+
+```text
+data/checkpoints/best.pth
+data/checkpoints/last.pth
+```
+
+`best.pth` stores the checkpoint with the lowest joint multi-task validation loss, while `last.pth` stores the final training epoch. Checkpoints generated from synthetic data validate the pipeline only and are not pretrained weights reproducing the paper results.
+
+### Training Weights
+
+This repository does not include synthetic, trained, or official pretrained weights. Running the training script generates `best.pth` and `last.pth` under `data/checkpoints/`. For real downstream use, train with appropriately converted SatlasPretrain data or use official AllenAI foundation model weights under their applicable license.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Prediction outputs:
+
+```text
+result/output/segmentation.npy
+result/output/regression.npy
+result/output/point.npy
+result/output/polygon.npy
+result/output/polyline.npy
+result/output/property.npy
+result/output/classification.npy
+```
+
+These files store class probabilities, dense regression values, confidence maps, and category probabilities for the seven compact task heads.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+Evaluation outputs:
+
+```text
+result/output/satlaspretrain_metrics.json
+result/output/qualitative_comparison.ppm
+```
+
+The metrics follow the task families reported in the paper: segmentation F1, regression MAE, point mAP, polygon mAP, polyline GEO, property accuracy, and classification accuracy. The qualitative comparison displays input imagery and paired ground-truth/prediction panels. Metrics and figures produced from synthetic data are pipeline checks, not paper results.
+
+# Official OneScience Resources
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| 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 |
+
+# Citation and License
+
+- This repository is a reproduction of the original SatlasPretrain paper.
+- The reproduction code is released under Apache License 2.0. The original SatlasPretrain data and official model weights remain subject to their respective source licenses.

conf/config.yaml

@@ -0,0 +1,21 @@
+model:
+  input_channels: 3
+  image_size: 64
+  time_steps: 4
+  embed_dim: 24
+  window_size: 4
+  segmentation_classes: 4
+  property_classes: 3
+  classification_classes: 3
+data:
+  path: ./data/fake_satlaspretrain.npz
+training:
+  epochs: 2
+  batch_size: 2
+  learning_rate: 0.0001
+  save_dir: ./data/checkpoints
+  checkpoint: ./data/checkpoints/best.pth
+runtime:
+  device: auto
+  seed: 42
+  output_dir: ./result/output

config.json

@@ -0,0 +1,63 @@
+{
+  "model_name": "SatlasPretrain",
+  "model_type": "satlaspretrain",
+  "architectures": [
+    "SatlasPretrain"
+  ],
+  "framework": "PyTorch",
+  "domain": "earth-science",
+  "task": "remote-sensing-image-understanding",
+  "implementation": {
+    "entry_point": "model/satlaspretrain.py",
+    "scope": "compact multi-temporal SatlasNet reproduction"
+  },
+  "architecture": {
+    "family": "multi-temporal multi-task SatlasNet",
+    "input_format": "BTCHW",
+    "input_channels": 3,
+    "image_size": 64,
+    "time_steps": 4,
+    "embed_dim": 24,
+    "window_size": 4,
+    "feature_scales": 4,
+    "temporal_aggregation": "max_pooling",
+    "task_heads": [
+      "segmentation",
+      "regression",
+      "point",
+      "polygon",
+      "polyline",
+      "property",
+      "classification"
+    ]
+  },
+  "data": {
+    "dataset": "SatlasPretrain",
+    "official_repo": "allenai/satlas-pretrain",
+    "image_modalities": [
+      "Sentinel-2",
+      "NAIP"
+    ],
+    "protocol": "synthetic-satlaspretrain-multitask",
+    "synthetic_input_shape": [
+      4,
+      4,
+      3,
+      64,
+      64
+    ]
+  },
+  "metrics": [
+    "segmentation_f1",
+    "regression_mae",
+    "point_map",
+    "polygon_map",
+    "polyline_geo",
+    "property_accuracy",
+    "classification_accuracy"
+  ],
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/satlaspretrain.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "remote_sensing_image_understanding",
+  "model": "SatlasPretrain",
+  "input_format": "BTCHW",
+  "protocol": "synthetic_satlaspretrain_multitask",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/satlaspretrain.py

@@ -0,0 +1,118 @@
+"""Compact SatlasNet reproduction for multi-temporal remote sensing data."""
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class WindowBlock(nn.Module):
+    """Window self-attention block used in the compact hierarchical backbone."""
+
+    def __init__(self, channels, window_size=4, heads=4):
+        super().__init__()
+        self.window_size = window_size
+        self.norm1 = nn.LayerNorm(channels)
+        self.attention = nn.MultiheadAttention(channels, heads, batch_first=True)
+        self.norm2 = nn.LayerNorm(channels)
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, channels * 4), nn.GELU(), nn.Linear(channels * 4, channels)
+        )
+
+    def forward(self, x):
+        batch, channels, height, width = x.shape
+        size = self.window_size
+        pad_h = (-height) % size
+        pad_w = (-width) % size
+        x = F.pad(x, (0, pad_w, 0, pad_h))
+        padded_h, padded_w = x.shape[-2:]
+        windows = x.permute(0, 2, 3, 1).reshape(
+            batch, padded_h // size, size, padded_w // size, size, channels
+        )
+        windows = windows.permute(0, 1, 3, 2, 4, 5).reshape(-1, size * size, channels)
+        normalized = self.norm1(windows)
+        attended = self.attention(normalized, normalized, normalized, need_weights=False)[0]
+        windows = windows + attended
+        windows = windows + self.mlp(self.norm2(windows))
+        x = windows.reshape(batch, padded_h // size, padded_w // size, size, size, channels)
+        x = x.permute(0, 1, 3, 2, 4, 5).reshape(batch, padded_h, padded_w, channels)
+        return x[:, :height, :width].permute(0, 3, 1, 2).contiguous()
+
+
+class SatlasPretrain(nn.Module):
+    """SatlasNet with hierarchical features, temporal max pooling, and seven heads."""
+
+    dense_tasks = ("segmentation", "regression", "point", "polygon", "polyline")
+
+    def __init__(
+        self,
+        input_channels=3,
+        embed_dim=24,
+        window_size=4,
+        segmentation_classes=4,
+        property_classes=3,
+        classification_classes=3,
+    ):
+        super().__init__()
+        dimensions = [embed_dim * (2**index) for index in range(4)]
+        self.patch_embed = nn.Conv2d(input_channels, dimensions[0], 4, stride=4)
+        self.stages = nn.ModuleList()
+        for index, channels in enumerate(dimensions):
+            downsample = nn.Identity() if index == 0 else nn.Conv2d(dimensions[index - 1], channels, 2, 2)
+            heads = max(1, channels // embed_dim)
+            self.stages.append(nn.Sequential(downsample, WindowBlock(channels, window_size, heads)))
+        self.lateral = nn.ModuleList(nn.Conv2d(channels, embed_dim, 1) for channels in dimensions)
+        dense_channels = {
+            "segmentation": segmentation_classes,
+            "regression": 1,
+            "point": 1,
+            "polygon": 1,
+            "polyline": 1,
+        }
+        self.dense_heads = nn.ModuleDict(
+            {name: nn.Conv2d(embed_dim, channels, 1) for name, channels in dense_channels.items()}
+        )
+        self.property_head = nn.Linear(dimensions[-1], property_classes)
+        self.classification_head = nn.Linear(dimensions[-1], classification_classes)
+
+    def encode_image(self, image):
+        features = []
+        x = self.patch_embed(image)
+        for stage in self.stages:
+            x = stage(x)
+            features.append(x)
+        return features
+
+    def forward(self, images):
+        if images.ndim != 5:
+            raise ValueError("images must have shape [batch, time, channels, height, width]")
+        batch, time, channels, height, width = images.shape
+        per_scale = self.encode_image(images.reshape(batch * time, channels, height, width))
+        pooled = [feature.reshape(batch, time, *feature.shape[1:]).amax(dim=1) for feature in per_scale]
+        pyramid = self.lateral[-1](pooled[-1])
+        for index in range(len(pooled) - 2, -1, -1):
+            pyramid = F.interpolate(pyramid, size=pooled[index].shape[-2:], mode="nearest")
+            pyramid = pyramid + self.lateral[index](pooled[index])
+        dense = {
+            name: F.interpolate(head(pyramid), size=(height, width), mode="bilinear", align_corners=False)
+            for name, head in self.dense_heads.items()
+        }
+        global_feature = pooled[-1].mean(dim=(-2, -1))
+        dense["property"] = self.property_head(global_feature)
+        dense["classification"] = self.classification_head(global_feature)
+        return dense
+
+
+def multitask_loss(outputs, targets):
+    losses = {
+        "segmentation": F.cross_entropy(outputs["segmentation"], targets["segmentation"].long()),
+        "regression": F.l1_loss(outputs["regression"], targets["regression"]),
+        "point": F.binary_cross_entropy_with_logits(outputs["point"], targets["point"]),
+        "polygon": F.binary_cross_entropy_with_logits(outputs["polygon"], targets["polygon"]),
+        "polyline": F.binary_cross_entropy_with_logits(outputs["polyline"], targets["polyline"]),
+        "property": F.cross_entropy(outputs["property"], targets["property"].long()),
+        "classification": F.cross_entropy(outputs["classification"], targets["classification"].long()),
+    }
+    return sum(losses.values()), losses
+
+
+__all__ = ["SatlasPretrain", "multitask_loss"]

scripts/fake_data.py

@@ -0,0 +1,49 @@
+import argparse
+from pathlib import Path
+
+import numpy as np
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate compact SatlasPretrain virtual data")
+    parser.add_argument("--output", default="data/fake_satlaspretrain.npz")
+    parser.add_argument("--samples", type=int, default=4)
+    parser.add_argument("--time-steps", type=int, default=4)
+    parser.add_argument("--image-size", type=int, default=64)
+    args = parser.parse_args()
+    rng = np.random.default_rng(42)
+    shape = (args.samples, args.image_size, args.image_size)
+    images = rng.random((args.samples, args.time_steps, 3, args.image_size, args.image_size), dtype=np.float32)
+    yy, xx = np.mgrid[: args.image_size, : args.image_size]
+    segmentation = np.empty(shape, dtype=np.int64)
+    point = np.zeros((args.samples, 1, args.image_size, args.image_size), dtype=np.float32)
+    polygon = np.zeros_like(point)
+    polyline = np.zeros_like(point)
+    for index in range(args.samples):
+        segmentation[index] = ((xx >= args.image_size // 2).astype(int) + 2 * (yy >= args.image_size // 2)).astype(np.int64)
+        center = args.image_size // 4 + index
+        point[index, 0, center - 1:center + 2, center - 1:center + 2] = 1
+        polygon[index, 0, args.image_size // 3:2 * args.image_size // 3, args.image_size // 3:2 * args.image_size // 3] = 1
+        polyline[index, 0, np.arange(args.image_size), (np.arange(args.image_size) + index) % args.image_size] = 1
+    regression = images[:, -1, :1] * 0.75 + images[:, 0, 1:2] * 0.25
+    properties = np.arange(args.samples, dtype=np.int64) % 3
+    classification = (images.mean(axis=(1, 2, 3, 4)) * 3).astype(np.int64).clip(0, 2)
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(
+        output,
+        images=images,
+        segmentation=segmentation,
+        regression=regression.astype(np.float32),
+        point=point,
+        polygon=polygon,
+        polyline=polyline,
+        property=properties,
+        classification=classification,
+    )
+    print(f"images: {images.shape}")
+    print(f"saved: {output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,52 @@
+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.satlaspretrain import SatlasPretrain
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run compact SatlasNet inference")
+    parser.add_argument("--config", default="conf/config.yaml")
+    parser.add_argument("--data")
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--output")
+    args = parser.parse_args()
+    with open(args.config, encoding="utf-8") as handle:
+        config = yaml.safe_load(handle)
+    model_config = config["model"]
+    model = SatlasPretrain(**{key: model_config[key] for key in (
+        "input_channels", "embed_dim", "window_size", "segmentation_classes",
+        "property_classes", "classification_classes"
+    )})
+    checkpoint = args.checkpoint or config["training"]["checkpoint"]
+    model.load_state_dict(torch.load(checkpoint, map_location="cpu", weights_only=True)["model"])
+    data = np.load(args.data or config["data"]["path"])
+    model.eval()
+    with torch.no_grad():
+        outputs = model(torch.from_numpy(data["images"]))
+    predictions = {}
+    for name, value in outputs.items():
+        if name == "segmentation":
+            predictions[name] = value.softmax(dim=1).numpy().astype(np.float32)
+        elif name in ("property", "classification"):
+            predictions[name] = value.softmax(dim=1).numpy().astype(np.float32)
+        elif name == "regression":
+            predictions[name] = value.numpy().astype(np.float32)
+        else:
+            predictions[name] = value.sigmoid().numpy().astype(np.float32)
+    output_dir = Path(args.output or config["runtime"]["output_dir"])
+    output_dir.mkdir(parents=True, exist_ok=True)
+    for task, prediction in predictions.items():
+        output = output_dir / f"{task}.npy"
+        np.save(output, prediction)
+        print(f"{task}: {prediction.shape} -> {output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,118 @@
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+
+def binary_f1(prediction, target, threshold=0.5):
+    prediction = prediction >= threshold
+    target = target >= 0.5
+    true_positive = np.logical_and(prediction, target).sum()
+    return float(2 * true_positive / (prediction.sum() + target.sum() + 1e-8))
+
+
+def best_f1(scores, target):
+    return max(binary_f1(scores, target, threshold) for threshold in np.arange(0.05, 1.0, 0.05))
+
+
+def average_precision(scores, target):
+    order = np.argsort(-scores.reshape(-1))
+    labels = target.reshape(-1)[order] >= 0.5
+    positives = labels.sum()
+    if positives == 0:
+        return 0.0
+    precision = np.cumsum(labels) / np.arange(1, labels.size + 1)
+    return float((precision * labels).sum() / positives)
+
+
+def dilate(mask, radius=2):
+    padded = np.pad(mask, ((0, 0), (0, 0), (radius, radius), (radius, radius)))
+    neighborhoods = []
+    height, width = mask.shape[-2:]
+    for row_shift in range(2 * radius + 1):
+        for col_shift in range(2 * radius + 1):
+            neighborhoods.append(padded[..., row_shift:row_shift + height, col_shift:col_shift + width])
+    return np.logical_or.reduce(neighborhoods)
+
+
+def geo_accuracy(scores, target):
+    prediction = scores >= 0.5
+    target = target >= 0.5
+    precision = np.logical_and(prediction, dilate(target)).sum() / (prediction.sum() + 1e-8)
+    recall = np.logical_and(target, dilate(prediction)).sum() / (target.sum() + 1e-8)
+    return float(2 * precision * recall / (precision + recall + 1e-8))
+
+
+def colorize_segmentation(segmentation):
+    palette = np.asarray([[0, 0, 0], [32, 180, 96], [40, 120, 220], [230, 180, 40]], dtype=np.uint8)
+    return palette[segmentation % len(palette)]
+
+
+def colorize_scalar(values):
+    values = np.clip(values, 0, 1)
+    return np.stack((values * 255, (1 - np.abs(values - 0.5) * 2) * 255, (1 - values) * 255), axis=-1).astype(np.uint8)
+
+
+def overlay(image, mask, color):
+    output = image.copy()
+    selected = mask >= 0.5
+    output[selected] = (0.35 * output[selected] + 0.65 * np.asarray(color)).astype(np.uint8)
+    return output
+
+
+def write_qualitative_comparison(path, target, prediction):
+    image = np.clip(target["images"][0, -1].transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)
+    panels = [
+        image,
+        colorize_segmentation(target["segmentation"][0]),
+        colorize_segmentation(prediction["segmentation"][0].argmax(axis=0)),
+        colorize_scalar(target["regression"][0, 0]),
+        colorize_scalar(prediction["regression"][0, 0]),
+        overlay(image, target["point"][0, 0], [255, 230, 40]),
+        overlay(image, prediction["point"][0, 0], [255, 230, 40]),
+        overlay(image, target["polygon"][0, 0], [255, 70, 190]),
+        overlay(image, prediction["polygon"][0, 0], [255, 70, 190]),
+        overlay(image, target["polyline"][0, 0], [40, 240, 255]),
+        overlay(image, prediction["polyline"][0, 0], [40, 240, 255]),
+    ]
+    separator = np.full((image.shape[0], 2, 3), 255, dtype=np.uint8)
+    comparison = panels[0]
+    for panel in panels[1:]:
+        comparison = np.concatenate((comparison, separator, panel), axis=1)
+    with open(path, "wb") as handle:
+        handle.write(f"P6\n{comparison.shape[1]} {comparison.shape[0]}\n255\n".encode("ascii"))
+        handle.write(comparison.tobytes())
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate and visualize SatlasNet predictions")
+    parser.add_argument("--prediction-dir", default="result/output")
+    parser.add_argument("--target", default="data/fake_satlaspretrain.npz")
+    args = parser.parse_args()
+    prediction_dir = Path(args.prediction_dir)
+    tasks = ("segmentation", "regression", "point", "polygon", "polyline", "property", "classification")
+    prediction = {task: np.load(prediction_dir / f"{task}.npy") for task in tasks}
+    target = np.load(args.target)
+    metrics = {
+        "Seg_F1": float(np.mean([best_f1(prediction["segmentation"][:, label], target["segmentation"] == label) for label in range(4)])),
+        "Reg_MAE": float(np.abs(prediction["regression"] - target["regression"]).mean()),
+        "Pt_mAP": average_precision(prediction["point"], target["point"]),
+        "Pgon_mAP": average_precision(prediction["polygon"], target["polygon"]),
+        "Pline_GEO": geo_accuracy(prediction["polyline"], target["polyline"]),
+        "Prop_Accuracy": float((prediction["property"].argmax(1) == target["property"]).mean()),
+        "Cls_Accuracy": float((prediction["classification"].argmax(1) == target["classification"]).mean()),
+    }
+    metrics_path = prediction_dir / "satlaspretrain_metrics.json"
+    with open(metrics_path, "w", encoding="utf-8") as handle:
+        json.dump({"image_mode": "low-resolution virtual data", "metrics": metrics}, handle, indent=2)
+    visualization_path = prediction_dir / "qualitative_comparison.ppm"
+    write_qualitative_comparison(visualization_path, target, prediction)
+    for key, value in metrics.items():
+        print(f"{key}={value:.6f}")
+    print(f"saved: {metrics_path}")
+    print(f"saved: {visualization_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,107 @@
+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.satlaspretrain import SatlasPretrain, multitask_loss
+
+
+class NpzDataset(Dataset):
+    def __init__(self, path):
+        data = np.load(path)
+        self.data = {key: data[key] for key in data.files}
+
+    def __len__(self):
+        return len(self.data["images"])
+
+    def __getitem__(self, index):
+        return {key: torch.as_tensor(value[index]) for key, value in self.data.items()}
+
+
+def load_config(path):
+    with open(path, encoding="utf-8") as handle:
+        return yaml.safe_load(handle)
+
+
+def build_model(config):
+    return SatlasPretrain(**{key: config["model"][key] for key in (
+        "input_channels", "embed_dim", "window_size", "segmentation_classes",
+        "property_classes", "classification_classes"
+    )})
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train compact SatlasNet")
+    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"]
+    device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() and requested != "cpu" else "cpu")
+    torch.manual_seed(config["runtime"]["seed"])
+    dataset = NpzDataset(args.data or config["data"]["path"])
+    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
+    loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], shuffle=sampler is None, sampler=sampler)
+    model = build_model(config).to(device)
+    if distributed:
+        model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
+    optimizer = torch.optim.Adam(model.parameters(), lr=config["training"]["learning_rate"])
+    train_history = []
+    val_history = []
+    best_validation = float("inf")
+    save_dir = Path(config["training"]["save_dir"])
+    epochs = args.epochs or config["training"]["epochs"]
+    for epoch in range(epochs):
+        if sampler is not None:
+            sampler.set_epoch(epoch)
+        model.train()
+        epoch_losses = []
+        for batch in loader:
+            batch = {key: value.to(device) for key, value in batch.items()}
+            total, _ = multitask_loss(model(batch["images"]), batch)
+            optimizer.zero_grad()
+            total.backward()
+            optimizer.step()
+            epoch_losses.append(total.item())
+        train_history.append(float(np.mean(epoch_losses)))
+        model.eval()
+        with torch.no_grad():
+            sample = {key: value[:1].to(device) for key, value in next(iter(loader)).items()}
+            validation, _ = multitask_loss(model(sample["images"]), sample)
+        val_history.append(validation.item())
+        if local_rank == 0:
+            print(f"epoch={epoch + 1} train_loss={train_history[-1]:.6f} val_loss={val_history[-1]:.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,
+                "train_loss": train_history[-1],
+                "validation_loss": val_history[-1],
+                "format": "satlaspretrain-compact-v1",
+            }
+            torch.save(checkpoint, save_dir / "last.pth")
+            if val_history[-1] < best_validation:
+                best_validation = val_history[-1]
+                torch.save(checkpoint, save_dir / "best.pth")
+    if local_rank == 0:
+        print(f"saved: {save_dir / 'last.pth'}")
+        print(f"best: {save_dir / 'best.pth'}")
+    if distributed:
+        torch.distributed.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()