Add English Hugging Face model card and configuration
Browse files- .gitattributes +5 -27
- README.md +177 -0
- conf/config.yaml +32 -0
- config.json +45 -0
- configuration.json +12 -0
- model/skysense.py +107 -0
- scripts/fake_data.py +66 -0
- scripts/inference.py +69 -0
- scripts/result.py +64 -0
- scripts/train.py +147 -0
- weight/.gitkeep +0 -0
.gitattributes
CHANGED
|
@@ -1,35 +1,13 @@
|
|
| 1 |
-
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
-
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
-
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
-
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
-
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
-
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
-
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
-
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
-
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
-
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
-
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
-
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
-
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
-
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
-
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
-
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
-
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
-
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
-
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
-
|
| 27 |
-
*.
|
|
|
|
|
|
|
|
|
|
| 28 |
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
-
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
-
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
-
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
-
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
-
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
-
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
| 1 |
*.bin filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
| 2 |
*.h5 filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 3 |
*.onnx filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 4 |
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 5 |
*.pth filter=lfs diff=lfs merge=lfs -text
|
|
|
|
| 6 |
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 12 |
*.tar filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
|
|
|
|
|
|
|
| 13 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
README.md
ADDED
|
@@ -0,0 +1,177 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
frameworks: PyTorch
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
license: apache-2.0
|
| 6 |
+
tags:
|
| 7 |
+
- OneScience
|
| 8 |
+
- Earth Science
|
| 9 |
+
- Remote Sensing
|
| 10 |
+
- Multi-Modal Remote Sensing
|
| 11 |
+
- Semantic Segmentation
|
| 12 |
+
- SkySense
|
| 13 |
+
- arxiv:2312.10115
|
| 14 |
+
tasks: []
|
| 15 |
+
datasets: []
|
| 16 |
+
---
|
| 17 |
+
<p align="center">
|
| 18 |
+
<strong>
|
| 19 |
+
<span style="font-size: 30px;">SkySense</span>
|
| 20 |
+
</strong>
|
| 21 |
+
</p>
|
| 22 |
+
|
| 23 |
+
# Model Introduction
|
| 24 |
+
|
| 25 |
+
SkySense is a multi-modal and temporal remote sensing foundation model for universal Earth observation interpretation. It jointly models high-resolution optical imagery, Sentinel-2 multispectral time series, Sentinel-1 SAR time series, and geographic context to provide unified representations for remote sensing classification, segmentation, detection, and change detection.
|
| 26 |
+
|
| 27 |
+
Paper: SkySense: A Multi-Modal Remote Sensing Foundation Model Towards Universal Interpretation for Earth Observation Imagery
|
| 28 |
+
|
| 29 |
+
https://arxiv.org/abs/2312.10115
|
| 30 |
+
|
| 31 |
+
# Model Description
|
| 32 |
+
|
| 33 |
+
SkySense was proposed by Ant Group, Wuhan University, and MYbank. The original model was trained on 21.5 million groups of high-resolution optical, Sentinel-2 multispectral, and Sentinel-1 SAR temporal samples, and is suitable for single-modal or multi-modal, static or temporal remote sensing classification and localization tasks.
|
| 34 |
+
|
| 35 |
+
# Use Cases
|
| 36 |
+
|
| 37 |
+
| Scenario | Description |
|
| 38 |
+
| :---: | :--- |
|
| 39 |
+
| Multi-modal remote sensing representation training | Train SkySense with multi-modal temporal remote sensing image data. |
|
| 40 |
+
| Remote sensing image semantic segmentation | Fuse optical, multispectral, and SAR features to predict land-cover classes for image pixels. |
|
| 41 |
+
| Local quick validation | Use synthetic data to validate data loading, training, inference, evaluation, and visualization. |
|
| 42 |
+
| Hugging Face / OneCode execution | Download the standalone model package, install dependencies, and run the scripts directly. |
|
| 43 |
+
| Multi-GPU training | Launch distributed multi-process training with `torchrun`. |
|
| 44 |
+
|
| 45 |
+
# Usage Guide
|
| 46 |
+
|
| 47 |
+
## 1. OneCode Usage
|
| 48 |
+
|
| 49 |
+
Experience intelligent one-click AI4S programming through the OneCode online environment:
|
| 50 |
+
|
| 51 |
+
[Click to Experience Intelligent One-Click AI4S Programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
|
| 52 |
+
|
| 53 |
+
## 2. Manual Installation and Usage
|
| 54 |
+
|
| 55 |
+
**Hardware Requirements**
|
| 56 |
+
|
| 57 |
+
- A GPU or DCU is recommended.
|
| 58 |
+
- CPU can be used for import and small-scale connectivity verification; full training and inference will be slower.
|
| 59 |
+
- DCU users must install DTK in advance. DTK 25.04.2 or above, or the OneScience recommended version matching the cluster, is recommended.
|
| 60 |
+
|
| 61 |
+
### Download the Model Package
|
| 62 |
+
|
| 63 |
+
```bash
|
| 64 |
+
hf download OneScience-Group/SkySense --local-dir ./SkySense
|
| 65 |
+
cd SkySense
|
| 66 |
+
```
|
| 67 |
+
|
| 68 |
+
### Install the Runtime Environment
|
| 69 |
+
|
| 70 |
+
**DCU Environment**
|
| 71 |
+
|
| 72 |
+
```bash
|
| 73 |
+
# Please activate DTK and CONDA first
|
| 74 |
+
conda create -n onescience311 python=3.11 -y
|
| 75 |
+
conda activate onescience311
|
| 76 |
+
# uv installation is supported
|
| 77 |
+
pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 78 |
+
```
|
| 79 |
+
|
| 80 |
+
**GPU Environment**
|
| 81 |
+
|
| 82 |
+
```bash
|
| 83 |
+
# Please activate CONDA first
|
| 84 |
+
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
|
| 85 |
+
conda activate onescience311
|
| 86 |
+
# uv installation is supported
|
| 87 |
+
pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
|
| 88 |
+
```
|
| 89 |
+
|
| 90 |
+
### Training Data Introduction
|
| 91 |
+
|
| 92 |
+
The original SkySense pretraining collection contains about 21.5 million spatially registered multi-modal remote sensing samples. It includes high-resolution optical imagery, Sentinel-2 multispectral temporal imagery, Sentinel-1 SAR temporal imagery, acquisition dates, geographic region information, and downstream task labels. The complete registered pretraining collection is not publicly available from the authors, and it is not included in this model repository.
|
| 93 |
+
|
| 94 |
+
Original data sources include Sentinel-1 and Sentinel-2 products available through Copernicus Data Space and commercial high-resolution imagery such as WorldView. Users must follow the applicable data licenses and independently complete cloud filtering, radiometric processing, spatial registration, resampling, cropping, normalization, and label preparation.
|
| 95 |
+
|
| 96 |
+
The default local validation uses synthetic NPZ data:
|
| 97 |
+
|
| 98 |
+
```bash
|
| 99 |
+
python scripts/fake_data.py
|
| 100 |
+
```
|
| 101 |
+
|
| 102 |
+
The command creates 24 training samples and 8 test samples. Synthetic data is used only to validate the complete training, inference, evaluation, and visualization pipeline. Synthetic results are not reported as model performance and do not reproduce the paper experiments.
|
| 103 |
+
|
| 104 |
+
For real-data training, prepare the corresponding multi-modal samples in the format expected by the scripts, save them under `data/`, update the data and model settings in `conf/config.yaml`, and use the same `scripts/train.py`. Do not run `scripts/fake_data.py` when using real data.
|
| 105 |
+
|
| 106 |
+
### Training
|
| 107 |
+
|
| 108 |
+
Synthetic-data quick validation:
|
| 109 |
+
|
| 110 |
+
```bash
|
| 111 |
+
python scripts/fake_data.py
|
| 112 |
+
python scripts/train.py
|
| 113 |
+
```
|
| 114 |
+
|
| 115 |
+
Single GPU or CPU:
|
| 116 |
+
|
| 117 |
+
```bash
|
| 118 |
+
python scripts/train.py
|
| 119 |
+
```
|
| 120 |
+
|
| 121 |
+
Multi-GPU:
|
| 122 |
+
|
| 123 |
+
```bash
|
| 124 |
+
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
|
| 125 |
+
```
|
| 126 |
+
|
| 127 |
+
Training outputs:
|
| 128 |
+
|
| 129 |
+
```text
|
| 130 |
+
result/checkpoints/skysense.pt
|
| 131 |
+
result/training/metrics.json
|
| 132 |
+
```
|
| 133 |
+
|
| 134 |
+
### Training Weights
|
| 135 |
+
|
| 136 |
+
Pretrained weights based on multi-modal temporal remote sensing data will be provided under `weight/`. The weight file will be uploaded in a future update. Checkpoints generated from synthetic data are only for pipeline validation and do not provide real remote sensing interpretation capability.
|
| 137 |
+
|
| 138 |
+
### Inference
|
| 139 |
+
|
| 140 |
+
```bash
|
| 141 |
+
python scripts/inference.py
|
| 142 |
+
```
|
| 143 |
+
|
| 144 |
+
Inference reads `data/test.npz` and `result/checkpoints/skysense.pt`. Test data must use the same channel order, normalization, spatial size, region mapping, and class definition as the training data.
|
| 145 |
+
|
| 146 |
+
Prediction output:
|
| 147 |
+
|
| 148 |
+
```text
|
| 149 |
+
result/output/
|
| 150 |
+
```
|
| 151 |
+
|
| 152 |
+
### Evaluation and Visualization
|
| 153 |
+
|
| 154 |
+
```bash
|
| 155 |
+
python scripts/result.py
|
| 156 |
+
```
|
| 157 |
+
|
| 158 |
+
Evaluation and visualization outputs:
|
| 159 |
+
|
| 160 |
+
```text
|
| 161 |
+
result/evaluation/metrics.json
|
| 162 |
+
result/evaluation/comparison.png
|
| 163 |
+
```
|
| 164 |
+
|
| 165 |
+
For synthetic data, evaluation outputs only confirm that the evaluation and visualization pipeline can run successfully. They are not reported as model performance. For real data, the metrics represent supervised semantic segmentation performance on the user-provided test set and are not equivalent to the complete 21.5-million-sample pretraining evaluation in the paper.
|
| 166 |
+
|
| 167 |
+
# Official OneScience Resources
|
| 168 |
+
|
| 169 |
+
| Platform | OneScience Main Repository | Skills Repository |
|
| 170 |
+
| --- | --- | --- |
|
| 171 |
+
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
|
| 172 |
+
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
|
| 173 |
+
|
| 174 |
+
# Citation and License
|
| 175 |
+
|
| 176 |
+
- This repository is a reproduction of the original SkySense paper.
|
| 177 |
+
- The reproduction code is released under Apache License 2.0. The original datasets and official model weights remain subject to their respective source licenses.
|
conf/config.yaml
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
seed: 42
|
| 2 |
+
data:
|
| 3 |
+
root: data
|
| 4 |
+
protocol: synthetic_skysense_npz
|
| 5 |
+
train_samples: 24
|
| 6 |
+
test_samples: 8
|
| 7 |
+
image_size: 32
|
| 8 |
+
num_classes: 6
|
| 9 |
+
hr_channels: 3
|
| 10 |
+
s2_channels: 10
|
| 11 |
+
s1_channels: 2
|
| 12 |
+
s2_timesteps: 4
|
| 13 |
+
s1_timesteps: 3
|
| 14 |
+
model:
|
| 15 |
+
embed_dim: 32
|
| 16 |
+
patch_size: 4
|
| 17 |
+
temporal_depth: 2
|
| 18 |
+
temporal_heads: 4
|
| 19 |
+
num_regions: 16
|
| 20 |
+
prototypes_per_region: 4
|
| 21 |
+
train:
|
| 22 |
+
epochs: 3
|
| 23 |
+
batch_size: 4
|
| 24 |
+
learning_rate: 0.001
|
| 25 |
+
weight_decay: 0.0001
|
| 26 |
+
alignment_weight: 0.05
|
| 27 |
+
num_workers: 0
|
| 28 |
+
paths:
|
| 29 |
+
checkpoint: result/checkpoints/skysense.pt
|
| 30 |
+
training_metrics: result/training/metrics.json
|
| 31 |
+
inference_dir: result/output
|
| 32 |
+
evaluation_dir: result/evaluation
|
config.json
ADDED
|
@@ -0,0 +1,45 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"model_name": "SkySense",
|
| 3 |
+
"model_type": "skysense",
|
| 4 |
+
"architectures": [
|
| 5 |
+
"SkySense"
|
| 6 |
+
],
|
| 7 |
+
"framework": "PyTorch",
|
| 8 |
+
"domain": "earth-science",
|
| 9 |
+
"task": "remote-sensing-multimodal-semantic-segmentation",
|
| 10 |
+
"implementation": {
|
| 11 |
+
"entry_point": "model/skysense.py",
|
| 12 |
+
"scope": "compact multi-modal temporal SkySense reproduction"
|
| 13 |
+
},
|
| 14 |
+
"architecture": {
|
| 15 |
+
"family": "multi-modal temporal remote sensing encoder",
|
| 16 |
+
"input_format": "NPZ multimodal temporal arrays",
|
| 17 |
+
"modalities": [
|
| 18 |
+
"high-resolution RGB",
|
| 19 |
+
"Sentinel-2 multispectral",
|
| 20 |
+
"Sentinel-1 SAR"
|
| 21 |
+
],
|
| 22 |
+
"image_size": 32,
|
| 23 |
+
"embed_dim": 32,
|
| 24 |
+
"patch_size": 4,
|
| 25 |
+
"temporal_depth": 2,
|
| 26 |
+
"temporal_heads": 4,
|
| 27 |
+
"task": "semantic segmentation"
|
| 28 |
+
},
|
| 29 |
+
"data": {
|
| 30 |
+
"dataset": "SkySense pretraining data",
|
| 31 |
+
"official_paper": "https://arxiv.org/abs/2312.10115",
|
| 32 |
+
"protocol": "synthetic_skysense_npz",
|
| 33 |
+
"synthetic_train_samples": 24,
|
| 34 |
+
"synthetic_test_samples": 8
|
| 35 |
+
},
|
| 36 |
+
"metrics": [
|
| 37 |
+
"pixel_accuracy",
|
| 38 |
+
"mean_iou",
|
| 39 |
+
"per_class_iou"
|
| 40 |
+
],
|
| 41 |
+
"configuration_sources": [
|
| 42 |
+
"conf/config.yaml",
|
| 43 |
+
"model/skysense.py"
|
| 44 |
+
]
|
| 45 |
+
}
|
configuration.json
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"framework": "PyTorch",
|
| 3 |
+
"task": "remote_sensing_semantic_segmentation",
|
| 4 |
+
"model": "SkySense",
|
| 5 |
+
"input_format": "multimodal_npz",
|
| 6 |
+
"protocol": "synthetic_skysense_npz",
|
| 7 |
+
"default_config": "conf/config.yaml",
|
| 8 |
+
"train": "scripts/train.py",
|
| 9 |
+
"inference": "scripts/inference.py",
|
| 10 |
+
"evaluation": "scripts/result.py",
|
| 11 |
+
"visualization": "scripts/result.py"
|
| 12 |
+
}
|
model/skysense.py
ADDED
|
@@ -0,0 +1,107 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Compact, trainable SkySense reproduction for multi-modal remote sensing data."""
|
| 2 |
+
|
| 3 |
+
import torch
|
| 4 |
+
from torch import nn
|
| 5 |
+
from torch.nn import functional as F
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
class SpatialEncoder(nn.Module):
|
| 9 |
+
def __init__(self, in_channels, embed_dim, patch_size):
|
| 10 |
+
super().__init__()
|
| 11 |
+
self.projection = nn.Sequential(
|
| 12 |
+
nn.Conv2d(in_channels, embed_dim, patch_size, patch_size),
|
| 13 |
+
nn.GELU(),
|
| 14 |
+
nn.Conv2d(embed_dim, embed_dim, 3, padding=1),
|
| 15 |
+
nn.GELU(),
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
def forward(self, images):
|
| 19 |
+
batch, time, channels, height, width = images.shape
|
| 20 |
+
features = self.projection(images.reshape(batch * time, channels, height, width))
|
| 21 |
+
_, dim, out_height, out_width = features.shape
|
| 22 |
+
return features.reshape(batch, time, dim, out_height, out_width)
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
class SkySense(nn.Module):
|
| 26 |
+
"""Factorized spatial-temporal encoder with geo-context prototypes."""
|
| 27 |
+
|
| 28 |
+
def __init__(
|
| 29 |
+
self,
|
| 30 |
+
hr_channels=3,
|
| 31 |
+
s2_channels=10,
|
| 32 |
+
s1_channels=2,
|
| 33 |
+
embed_dim=32,
|
| 34 |
+
patch_size=4,
|
| 35 |
+
temporal_depth=2,
|
| 36 |
+
temporal_heads=4,
|
| 37 |
+
num_regions=16,
|
| 38 |
+
prototypes_per_region=4,
|
| 39 |
+
num_classes=6,
|
| 40 |
+
):
|
| 41 |
+
super().__init__()
|
| 42 |
+
self.patch_size = patch_size
|
| 43 |
+
self.num_regions = num_regions
|
| 44 |
+
self.hr_encoder = SpatialEncoder(hr_channels, embed_dim, patch_size)
|
| 45 |
+
self.s2_encoder = SpatialEncoder(s2_channels, embed_dim, patch_size)
|
| 46 |
+
self.s1_encoder = SpatialEncoder(s1_channels, embed_dim, patch_size)
|
| 47 |
+
self.date_embedding = nn.Embedding(366, embed_dim)
|
| 48 |
+
self.modality_embedding = nn.Parameter(torch.zeros(3, embed_dim))
|
| 49 |
+
self.fusion_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
|
| 50 |
+
layer = nn.TransformerEncoderLayer(
|
| 51 |
+
d_model=embed_dim,
|
| 52 |
+
nhead=temporal_heads,
|
| 53 |
+
dim_feedforward=embed_dim * 4,
|
| 54 |
+
dropout=0.0,
|
| 55 |
+
activation="gelu",
|
| 56 |
+
batch_first=True,
|
| 57 |
+
norm_first=True,
|
| 58 |
+
)
|
| 59 |
+
self.temporal_fusion = nn.TransformerEncoder(layer, temporal_depth)
|
| 60 |
+
self.prototypes = nn.Parameter(
|
| 61 |
+
torch.randn(num_regions, prototypes_per_region, embed_dim) * 0.02
|
| 62 |
+
)
|
| 63 |
+
self.decoder = nn.Sequential(
|
| 64 |
+
nn.Conv2d(embed_dim * 2, embed_dim, 3, padding=1),
|
| 65 |
+
nn.GELU(),
|
| 66 |
+
nn.Conv2d(embed_dim, num_classes, 1),
|
| 67 |
+
)
|
| 68 |
+
nn.init.normal_(self.date_embedding.weight, std=0.02)
|
| 69 |
+
nn.init.normal_(self.modality_embedding, std=0.02)
|
| 70 |
+
nn.init.normal_(self.fusion_token, std=0.02)
|
| 71 |
+
|
| 72 |
+
def _add_context(self, features, dates, modality_index):
|
| 73 |
+
date_context = self.date_embedding(dates.clamp(0, 365)).unsqueeze(-1).unsqueeze(-1)
|
| 74 |
+
modality = self.modality_embedding[modality_index].view(1, 1, -1, 1, 1)
|
| 75 |
+
return features + date_context + modality
|
| 76 |
+
|
| 77 |
+
def encode_modalities(self, hr, s2, s1, dates_hr, dates_s2, dates_s1):
|
| 78 |
+
return (
|
| 79 |
+
self._add_context(self.hr_encoder(hr), dates_hr, 0),
|
| 80 |
+
self._add_context(self.s2_encoder(s2), dates_s2, 1),
|
| 81 |
+
self._add_context(self.s1_encoder(s1), dates_s1, 2),
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
def forward(self, hr, s2, s1, dates_hr, dates_s2, dates_s1, region):
|
| 85 |
+
modality_features = self.encode_modalities(hr, s2, s1, dates_hr, dates_s2, dates_s1)
|
| 86 |
+
temporal = torch.cat(modality_features, dim=1)
|
| 87 |
+
batch, time, dim, out_height, out_width = temporal.shape
|
| 88 |
+
temporal = temporal.permute(0, 3, 4, 1, 2).reshape(-1, time, dim)
|
| 89 |
+
token = self.fusion_token.expand(temporal.shape[0], -1, -1)
|
| 90 |
+
fused = self.temporal_fusion(torch.cat([token, temporal], dim=1))[:, 0]
|
| 91 |
+
fused = fused.reshape(batch, out_height, out_width, dim)
|
| 92 |
+
|
| 93 |
+
regional_prototypes = self.prototypes[region.remainder(self.num_regions)]
|
| 94 |
+
query = F.normalize(fused, dim=-1)
|
| 95 |
+
keys = F.normalize(regional_prototypes, dim=-1)
|
| 96 |
+
attention = torch.einsum("bhwd,bpd->bhwp", query, keys).softmax(dim=-1)
|
| 97 |
+
geo_context = torch.einsum("bhwp,bpd->bhwd", attention, regional_prototypes)
|
| 98 |
+
output = torch.cat([fused, geo_context], dim=-1).permute(0, 3, 1, 2)
|
| 99 |
+
logits = self.decoder(output)
|
| 100 |
+
logits = F.interpolate(logits, size=hr.shape[-2:], mode="bilinear", align_corners=False)
|
| 101 |
+
return {"logits": logits, "features": modality_features, "fused": fused}
|
| 102 |
+
|
| 103 |
+
@staticmethod
|
| 104 |
+
def cross_modal_alignment_loss(features):
|
| 105 |
+
pooled = [F.normalize(feature.mean(dim=(1, 3, 4)), dim=-1) for feature in features]
|
| 106 |
+
losses = [1.0 - (pooled[i] * pooled[j]).sum(dim=-1).mean() for i in range(3) for j in range(i + 1, 3)]
|
| 107 |
+
return torch.stack(losses).mean()
|
scripts/fake_data.py
ADDED
|
@@ -0,0 +1,66 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Generate deterministic SkySense-format data for connectivity tests."""
|
| 2 |
+
|
| 3 |
+
import argparse
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import yaml
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def load_config():
|
| 14 |
+
with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
|
| 15 |
+
return yaml.safe_load(handle)
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def make_split(path, count, config, seed):
|
| 19 |
+
rng = np.random.default_rng(seed)
|
| 20 |
+
data = config["data"]
|
| 21 |
+
size = data["image_size"]
|
| 22 |
+
hr = rng.normal(size=(count, 1, data["hr_channels"], size, size)).astype("float32")
|
| 23 |
+
s2 = rng.normal(size=(count, data["s2_timesteps"], data["s2_channels"], size, size)).astype("float32")
|
| 24 |
+
s1 = rng.normal(size=(count, data["s1_timesteps"], data["s1_channels"], size, size)).astype("float32")
|
| 25 |
+
dates_hr = rng.integers(0, 365, size=(count, 1), dtype=np.int64)
|
| 26 |
+
dates_s2 = rng.integers(0, 365, size=(count, data["s2_timesteps"]), dtype=np.int64)
|
| 27 |
+
dates_s1 = rng.integers(0, 365, size=(count, data["s1_timesteps"]), dtype=np.int64)
|
| 28 |
+
region = rng.integers(0, config["model"]["num_regions"], size=count, dtype=np.int64)
|
| 29 |
+
signal = hr[:, 0, 0] + 0.35 * s2[:, :, 3].mean(axis=1) - 0.2 * s1[:, :, 0].mean(axis=1)
|
| 30 |
+
bins = np.quantile(signal, np.linspace(0, 1, data["num_classes"] + 1)[1:-1])
|
| 31 |
+
labels = np.digitize(signal, bins).astype(np.int64)
|
| 32 |
+
np.savez_compressed(
|
| 33 |
+
path,
|
| 34 |
+
hr=hr,
|
| 35 |
+
s2=s2,
|
| 36 |
+
s1=s1,
|
| 37 |
+
dates_hr=dates_hr,
|
| 38 |
+
dates_s2=dates_s2,
|
| 39 |
+
dates_s1=dates_s1,
|
| 40 |
+
region=region,
|
| 41 |
+
labels=labels,
|
| 42 |
+
data_source=np.asarray("synthetic"),
|
| 43 |
+
protocol=np.asarray(config["data"]["protocol"]),
|
| 44 |
+
)
|
| 45 |
+
|
| 46 |
+
|
| 47 |
+
def main():
|
| 48 |
+
parser = argparse.ArgumentParser()
|
| 49 |
+
parser.add_argument("--force", action="store_true")
|
| 50 |
+
args = parser.parse_args()
|
| 51 |
+
config = load_config()
|
| 52 |
+
data_dir = ROOT / config["data"]["root"]
|
| 53 |
+
data_dir.mkdir(parents=True, exist_ok=True)
|
| 54 |
+
outputs = [("train.npz", config["data"]["train_samples"]), ("test.npz", config["data"]["test_samples"])]
|
| 55 |
+
for offset, (name, count) in enumerate(outputs):
|
| 56 |
+
target = data_dir / name
|
| 57 |
+
if args.force or not target.exists():
|
| 58 |
+
make_split(target, count, config, config["seed"] + offset)
|
| 59 |
+
print(
|
| 60 |
+
f"generated={target.relative_to(ROOT)} samples={count} "
|
| 61 |
+
f"data_source=synthetic protocol={config['data']['protocol']}"
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
if __name__ == "__main__":
|
| 66 |
+
main()
|
scripts/inference.py
ADDED
|
@@ -0,0 +1,69 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Run SkySense inference and save arrays for evaluation."""
|
| 2 |
+
|
| 3 |
+
import importlib.util
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import numpy as np
|
| 7 |
+
import torch
|
| 8 |
+
import yaml
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def load_model_class():
|
| 15 |
+
spec = importlib.util.spec_from_file_location("skysense_model", ROOT / "model" / "skysense.py")
|
| 16 |
+
module = importlib.util.module_from_spec(spec)
|
| 17 |
+
spec.loader.exec_module(module)
|
| 18 |
+
return module.SkySense
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def main():
|
| 22 |
+
with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
|
| 23 |
+
config = yaml.safe_load(handle)
|
| 24 |
+
checkpoint_path = ROOT / config["paths"]["checkpoint"]
|
| 25 |
+
if not checkpoint_path.exists():
|
| 26 |
+
raise FileNotFoundError(
|
| 27 |
+
f"Missing checkpoint: {checkpoint_path.relative_to(ROOT)}. "
|
| 28 |
+
"Run `python scripts/train.py` first."
|
| 29 |
+
)
|
| 30 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 31 |
+
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
| 32 |
+
SkySense = load_model_class()
|
| 33 |
+
model = SkySense(
|
| 34 |
+
**config["model"],
|
| 35 |
+
hr_channels=config["data"]["hr_channels"],
|
| 36 |
+
s2_channels=config["data"]["s2_channels"],
|
| 37 |
+
s1_channels=config["data"]["s1_channels"],
|
| 38 |
+
num_classes=config["data"]["num_classes"],
|
| 39 |
+
).to(device)
|
| 40 |
+
model.load_state_dict(checkpoint["model"])
|
| 41 |
+
model.eval()
|
| 42 |
+
test_path = ROOT / config["data"]["root"] / "test.npz"
|
| 43 |
+
if not test_path.exists():
|
| 44 |
+
raise FileNotFoundError(
|
| 45 |
+
f"Missing inference data: {test_path.relative_to(ROOT)}. "
|
| 46 |
+
"Run `python scripts/fake_data.py` first."
|
| 47 |
+
)
|
| 48 |
+
archive = np.load(test_path)
|
| 49 |
+
keys = ["hr", "s2", "s1", "dates_hr", "dates_s2", "dates_s1", "region"]
|
| 50 |
+
tensors = {key: torch.from_numpy(archive[key]).to(device) for key in keys}
|
| 51 |
+
with torch.inference_mode():
|
| 52 |
+
output = model(tensors["hr"], tensors["s2"], tensors["s1"], tensors["dates_hr"], tensors["dates_s2"], tensors["dates_s1"], tensors["region"])
|
| 53 |
+
probabilities = output["logits"].softmax(dim=1)
|
| 54 |
+
output_dir = ROOT / config["paths"]["inference_dir"]
|
| 55 |
+
output_dir.mkdir(parents=True, exist_ok=True)
|
| 56 |
+
np.save(output_dir / "predictions.npy", probabilities.argmax(dim=1).cpu().numpy())
|
| 57 |
+
np.save(output_dir / "probabilities.npy", probabilities.cpu().numpy())
|
| 58 |
+
np.save(output_dir / "targets.npy", archive["labels"])
|
| 59 |
+
data_source = str(archive["data_source"]) if "data_source" in archive.files else "unknown"
|
| 60 |
+
protocol = str(archive["protocol"]) if "protocol" in archive.files else "unknown"
|
| 61 |
+
np.savez(output_dir / "metadata.npz", data_source=data_source, protocol=protocol)
|
| 62 |
+
print(
|
| 63 |
+
f"output={output_dir.relative_to(ROOT)} samples={len(archive['hr'])} "
|
| 64 |
+
f"data_source={data_source} protocol={protocol}"
|
| 65 |
+
)
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
if __name__ == "__main__":
|
| 69 |
+
main()
|
scripts/result.py
ADDED
|
@@ -0,0 +1,64 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Evaluate semantic predictions and render a compact comparison image."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
import matplotlib.pyplot as plt
|
| 7 |
+
import numpy as np
|
| 8 |
+
import yaml
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def main():
|
| 15 |
+
with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
|
| 16 |
+
config = yaml.safe_load(handle)
|
| 17 |
+
input_dir = ROOT / config["paths"]["inference_dir"]
|
| 18 |
+
required = [input_dir / "predictions.npy", input_dir / "targets.npy"]
|
| 19 |
+
missing = [str(path.relative_to(ROOT)) for path in required if not path.exists()]
|
| 20 |
+
if missing:
|
| 21 |
+
raise FileNotFoundError(
|
| 22 |
+
f"Missing inference outputs: {missing}. Run `python scripts/inference.py` first."
|
| 23 |
+
)
|
| 24 |
+
predictions = np.load(input_dir / "predictions.npy")
|
| 25 |
+
targets = np.load(input_dir / "targets.npy")
|
| 26 |
+
classes = config["data"]["num_classes"]
|
| 27 |
+
intersections = np.zeros(classes, dtype=np.float64)
|
| 28 |
+
unions = np.zeros(classes, dtype=np.float64)
|
| 29 |
+
for class_id in range(classes):
|
| 30 |
+
predicted = predictions == class_id
|
| 31 |
+
expected = targets == class_id
|
| 32 |
+
intersections[class_id] = np.logical_and(predicted, expected).sum()
|
| 33 |
+
unions[class_id] = np.logical_or(predicted, expected).sum()
|
| 34 |
+
per_class_iou = np.divide(intersections, unions, out=np.zeros_like(intersections), where=unions > 0)
|
| 35 |
+
metadata_path = input_dir / "metadata.npz"
|
| 36 |
+
metadata = np.load(metadata_path) if metadata_path.exists() else None
|
| 37 |
+
metrics = {
|
| 38 |
+
"pixel_accuracy": float((predictions == targets).mean()),
|
| 39 |
+
"mean_iou": float(per_class_iou.mean()),
|
| 40 |
+
"per_class_iou": per_class_iou.tolist(),
|
| 41 |
+
"samples": int(len(predictions)),
|
| 42 |
+
"data_source": str(metadata["data_source"]) if metadata is not None else "unknown",
|
| 43 |
+
"protocol": str(metadata["protocol"]) if metadata is not None else "unknown",
|
| 44 |
+
}
|
| 45 |
+
output_dir = ROOT / config["paths"]["evaluation_dir"]
|
| 46 |
+
output_dir.mkdir(parents=True, exist_ok=True)
|
| 47 |
+
with (output_dir / "metrics.json").open("w", encoding="utf-8") as handle:
|
| 48 |
+
json.dump(metrics, handle, indent=2)
|
| 49 |
+
figure, axes = plt.subplots(2, 2, figsize=(7, 7))
|
| 50 |
+
for index, axis in enumerate(axes.flat):
|
| 51 |
+
sample = index // 2
|
| 52 |
+
image = targets[sample] if index % 2 == 0 else predictions[sample]
|
| 53 |
+
axis.imshow(image, vmin=0, vmax=classes - 1, cmap="terrain")
|
| 54 |
+
axis.set_title(("Target" if index % 2 == 0 else "Prediction") + f" {sample}")
|
| 55 |
+
axis.axis("off")
|
| 56 |
+
figure.tight_layout()
|
| 57 |
+
figure.savefig(output_dir / "comparison.png", dpi=120)
|
| 58 |
+
plt.close(figure)
|
| 59 |
+
print(f"pixel_accuracy={metrics['pixel_accuracy']:.6f} mean_iou={metrics['mean_iou']:.6f}")
|
| 60 |
+
print(f"evaluation={output_dir.relative_to(ROOT)}")
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
if __name__ == "__main__":
|
| 64 |
+
main()
|
scripts/train.py
ADDED
|
@@ -0,0 +1,147 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Train compact SkySense on NPZ multi-modal temporal samples."""
|
| 2 |
+
|
| 3 |
+
import importlib.util
|
| 4 |
+
import json
|
| 5 |
+
import os
|
| 6 |
+
import random
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
|
| 9 |
+
import numpy as np
|
| 10 |
+
import torch
|
| 11 |
+
import yaml
|
| 12 |
+
from torch import distributed as dist
|
| 13 |
+
from torch.nn.parallel import DistributedDataParallel
|
| 14 |
+
from torch.utils.data import DataLoader, Dataset, DistributedSampler
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
ROOT = Path(__file__).resolve().parents[1]
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def load_model_class():
|
| 21 |
+
spec = importlib.util.spec_from_file_location("skysense_model", ROOT / "model" / "skysense.py")
|
| 22 |
+
module = importlib.util.module_from_spec(spec)
|
| 23 |
+
spec.loader.exec_module(module)
|
| 24 |
+
return module.SkySense
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def load_config():
|
| 28 |
+
with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle:
|
| 29 |
+
return yaml.safe_load(handle)
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
class NPZDataset(Dataset):
|
| 33 |
+
def __init__(self, path):
|
| 34 |
+
archive = np.load(path)
|
| 35 |
+
sample_keys = {"hr", "s2", "s1", "dates_hr", "dates_s2", "dates_s1", "region", "labels"}
|
| 36 |
+
missing = sample_keys.difference(archive.files)
|
| 37 |
+
if missing:
|
| 38 |
+
raise ValueError(f"Dataset {path} is missing arrays: {sorted(missing)}")
|
| 39 |
+
self.arrays = {key: archive[key] for key in sample_keys}
|
| 40 |
+
self.data_source = str(archive["data_source"]) if "data_source" in archive.files else "unknown"
|
| 41 |
+
self.protocol = str(archive["protocol"]) if "protocol" in archive.files else "unknown"
|
| 42 |
+
|
| 43 |
+
def __len__(self):
|
| 44 |
+
return len(self.arrays["hr"])
|
| 45 |
+
|
| 46 |
+
def __getitem__(self, index):
|
| 47 |
+
return {key: torch.as_tensor(value[index]) for key, value in self.arrays.items()}
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def setup_device():
|
| 51 |
+
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
| 52 |
+
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
| 53 |
+
if world_size > 1:
|
| 54 |
+
backend = "nccl" if torch.cuda.is_available() else "gloo"
|
| 55 |
+
dist.init_process_group(backend=backend)
|
| 56 |
+
if torch.cuda.is_available():
|
| 57 |
+
torch.cuda.set_device(local_rank)
|
| 58 |
+
return torch.device("cuda", local_rank), world_size, local_rank
|
| 59 |
+
return torch.device("cpu"), world_size, local_rank
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def main():
|
| 63 |
+
config = load_config()
|
| 64 |
+
seed = config["seed"]
|
| 65 |
+
random.seed(seed)
|
| 66 |
+
np.random.seed(seed)
|
| 67 |
+
torch.manual_seed(seed)
|
| 68 |
+
device, world_size, local_rank = setup_device()
|
| 69 |
+
dataset_path = ROOT / config["data"]["root"] / "train.npz"
|
| 70 |
+
if not dataset_path.exists():
|
| 71 |
+
raise FileNotFoundError(
|
| 72 |
+
f"Missing training data: {dataset_path.relative_to(ROOT)}. "
|
| 73 |
+
"Run `python scripts/fake_data.py` for a synthetic connectivity test."
|
| 74 |
+
)
|
| 75 |
+
dataset = NPZDataset(dataset_path)
|
| 76 |
+
if local_rank == 0:
|
| 77 |
+
print(
|
| 78 |
+
f"data_source={dataset.data_source} protocol={dataset.protocol} "
|
| 79 |
+
f"samples={len(dataset)} image_size={config['data']['image_size']}x{config['data']['image_size']}"
|
| 80 |
+
)
|
| 81 |
+
sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None
|
| 82 |
+
loader = DataLoader(
|
| 83 |
+
dataset,
|
| 84 |
+
batch_size=config["train"]["batch_size"],
|
| 85 |
+
shuffle=sampler is None,
|
| 86 |
+
sampler=sampler,
|
| 87 |
+
num_workers=config["train"]["num_workers"],
|
| 88 |
+
)
|
| 89 |
+
SkySense = load_model_class()
|
| 90 |
+
model = SkySense(
|
| 91 |
+
**config["model"],
|
| 92 |
+
hr_channels=config["data"]["hr_channels"],
|
| 93 |
+
s2_channels=config["data"]["s2_channels"],
|
| 94 |
+
s1_channels=config["data"]["s1_channels"],
|
| 95 |
+
num_classes=config["data"]["num_classes"],
|
| 96 |
+
).to(device)
|
| 97 |
+
if world_size > 1:
|
| 98 |
+
model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None)
|
| 99 |
+
optimizer = torch.optim.AdamW(model.parameters(), lr=config["train"]["learning_rate"], weight_decay=config["train"]["weight_decay"])
|
| 100 |
+
final_loss = float("nan")
|
| 101 |
+
for epoch in range(config["train"]["epochs"]):
|
| 102 |
+
if sampler is not None:
|
| 103 |
+
sampler.set_epoch(epoch)
|
| 104 |
+
model.train()
|
| 105 |
+
total = 0.0
|
| 106 |
+
for batch in loader:
|
| 107 |
+
batch = {key: value.to(device) for key, value in batch.items()}
|
| 108 |
+
output = model(batch["hr"], batch["s2"], batch["s1"], batch["dates_hr"], batch["dates_s2"], batch["dates_s1"], batch["region"])
|
| 109 |
+
base_model = model.module if hasattr(model, "module") else model
|
| 110 |
+
segmentation = torch.nn.functional.cross_entropy(output["logits"], batch["labels"])
|
| 111 |
+
alignment = base_model.cross_modal_alignment_loss(output["features"])
|
| 112 |
+
loss = segmentation + config["train"]["alignment_weight"] * alignment
|
| 113 |
+
optimizer.zero_grad(set_to_none=True)
|
| 114 |
+
loss.backward()
|
| 115 |
+
optimizer.step()
|
| 116 |
+
total += loss.item()
|
| 117 |
+
final_loss = total / len(loader)
|
| 118 |
+
if local_rank == 0:
|
| 119 |
+
print(f"epoch={epoch + 1} loss={final_loss:.6f}")
|
| 120 |
+
if local_rank == 0:
|
| 121 |
+
checkpoint = ROOT / config["paths"]["checkpoint"]
|
| 122 |
+
checkpoint.parent.mkdir(parents=True, exist_ok=True)
|
| 123 |
+
base_model = model.module if hasattr(model, "module") else model
|
| 124 |
+
torch.save({"model": base_model.state_dict(), "config": config, "final_loss": final_loss}, checkpoint)
|
| 125 |
+
metrics = ROOT / config["paths"]["training_metrics"]
|
| 126 |
+
metrics.parent.mkdir(parents=True, exist_ok=True)
|
| 127 |
+
metrics.write_text(
|
| 128 |
+
json.dumps(
|
| 129 |
+
{
|
| 130 |
+
"final_loss": final_loss,
|
| 131 |
+
"epochs": config["train"]["epochs"],
|
| 132 |
+
"samples": len(dataset),
|
| 133 |
+
"data_source": dataset.data_source,
|
| 134 |
+
"protocol": dataset.protocol,
|
| 135 |
+
},
|
| 136 |
+
indent=2,
|
| 137 |
+
)
|
| 138 |
+
+ "\n",
|
| 139 |
+
encoding="utf-8",
|
| 140 |
+
)
|
| 141 |
+
print(f"checkpoint={checkpoint.relative_to(ROOT)} final_loss={final_loss:.6f}")
|
| 142 |
+
if world_size > 1:
|
| 143 |
+
dist.destroy_process_group()
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
if __name__ == "__main__":
|
| 147 |
+
main()
|
weight/.gitkeep
ADDED
|
File without changes
|