PSRGAN v1.0

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* 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
+images/psrgan_city.png filter=lfs diff=lfs merge=lfs -text
+images/psrgan_uncertainty.png filter=lfs diff=lfs merge=lfs -text
+psrgan_grid.jpg filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,30 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Mark-Daniel Leupold
+Copyright (c) 2021, Xintao Wang (Real-ESRGAN, on whose weights and RRDBNet architecture PSRGAN is based)
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -0,0 +1,187 @@
+---
+license: bsd-3-clause
+library_name: pytorch
+pipeline_tag: image-to-image
+language:
+- en
+tags:
+- super-resolution
+- sentinel-2
+- remote-sensing
+- earth-observation
+- uncertainty
+- bayesian-deep-learning
+- swag
+- esrgan
+- gan
+datasets:
+- marklpd/PSRGAN-Sentinel-2
+base_model: xinntao/Real-ESRGAN
+metrics:
+- psnr
+- ssim
+- lpips
+model-index:
+- name: PSRGAN
+  results:
+  - task:
+      type: image-to-image
+      name: 4x super-resolution
+    dataset:
+      type: marklpd/PSRGAN-Sentinel-2
+      name: PSRGAN Sentinel-2 Super-Resolution (test split)
+      split: test
+    metrics:
+    - type: psnr
+      value: 21.28
+      name: PSNR (dB)
+    - type: ssim
+      value: 0.5694
+      name: SSIM
+    - type: lpips
+      value: 0.1879
+      name: LPIPS (AlexNet)
+---
+![PSRGAN grid](psrgan_grid.jpg)
+
+*Test patches from regions the model never saw during training. Error and uncertainty share one color scale per row.*
+
+`PSRGAN` is a probabilistic 4× super-resolution model for Sentinel-2 RGB imagery. Next to the sharpened image it predicts a per-pixel uncertainty map, so you can see where the model is guessing. It is a Bayesian Real-ESRGAN with 16.7 million parameters, trained on the Alpine region with [PSRGAN Sentinel-2](https://huggingface.co/datasets/marklpd/PSRGAN-Sentinel-2).
+
+# Key Features
+
+1. Every output comes with an uncertainty map that combines aleatoric uncertainty (missing information in the input) and epistemic uncertainty (uncertainty of the model itself, from SWAG weight samples).
+2. The uncertainty tracks the error. Pixel-wise uncertainty and absolute error correlate with Pearson r = 0.59 (Spearman ρ = 0.70), and the 10% most uncertain pixels have 14.3× the error of the 10% most confident ones.
+3. On 1,056 test patches from unseen regions, PSRGAN reaches an LPIPS of 0.188, against 0.348 for Real-ESRGAN x4plus and 0.599 for bicubic upsampling. It also beats Real-ESRGAN in PSNR (21.28 vs. 20.22 dB) and SSIM (0.569 vs. 0.475).
+4. Two sets of weights. The final adversarial weights have the lowest LPIPS, the weights after the likelihood stage the highest PSNR (21.94 dB) and SSIM (0.598) of all four methods tested.
+5. Trained on randomized degradations (Gaussian, motion or defocus blur, then bicubic, bilinear or area downsampling), so it learns super-resolution instead of inverting one fixed downscaling operation.
+
+# How the uncertainty map is made
+
+A 64×64 input does not contain the detail of a 256×256 image, so every super-resolution model has to guess. PSRGAN models each output pixel as a Gaussian instead of a single value, and the uncertainty map is its standard deviation σ.
+
+**Aleatoric uncertainty** comes from the data. The last convolution of Real-ESRGAN is replaced by two Bayesian 3×3 heads: the μ-head outputs the image, the s-head the log-variance s = log σ² of every pixel. A Gaussian negative log-likelihood loss punishes large errors and confident errors alike, so σ_aleatoric = exp(s / 2) grows where the input is ambiguous.
+
+**Epistemic uncertainty** comes from the model. During the last 8 adversarial epochs, SWAG stores 32 weight snapshots and fits a Gaussian over the weights. Sampling 16 weight sets from it gives 16 slightly different images, and their per-pixel spread is σ_epistemic. Both parts combine to σ_total = √(σ²_aleatoric + 10 · σ²_epistemic).
+
+<center>
+  <img src="images/psrgan_uncertainty.png" width="900px">
+  <br><em>From prediction to uncertainty map, each map on its own color scale.</em>
+</center>
+
+In the patches shown here, forests, water and shaded slopes stay dark in both the error and the uncertainty map, while ridges, snow edges, roads, buildings and field borders light up in both.
+
+# Usage
+
+The model code is a single file, [`psrgan.py`](psrgan.py), that only needs `torch`, `huggingface_hub` and `safetensors`.
+
+```shell
+pip install torch huggingface_hub safetensors
+```
+
+```python
+import torch
+from huggingface_hub import hf_hub_download
+
+hf_hub_download("marklpd/PSRGAN", "psrgan.py", local_dir=".")
+from psrgan import PSRGAN
+
+model = PSRGAN.from_pretrained("marklpd/PSRGAN").eval()
+model.load_swag(hf_hub_download("marklpd/PSRGAN", "psrgan_x4_swag.safetensors"))
+
+lr = torch.rand(1, 3, 64, 64)            # RGB (B04, B03, B02), stretched to [0, 1]
+out = model.super_resolve(lr, swag_samples=16)
+out["sr"]                                # (1, 3, 256, 256) super-resolved image
+out["sigma_total"]                       # (1, 3, 256, 256) uncertainty map
+out["sigma_aleatoric"], out["sigma_epistemic"]
+```
+
+Without `load_swag`, `model.super_resolve(lr)` returns the image and the aleatoric map in one forward pass. `mc_samples=8` samples the Bayesian heads instead of SWAG, which is cheaper but only covers the last layer.
+
+The model expects Sentinel-2 L2A RGB (B04, B03, B02) after a per-image 1 to 99% percentile stretch, the same preprocessing as in the dataset. `psrgan.stretch(reflectance)` applies it. The network is fully convolutional, so any input size works. Large scenes should be processed in overlapping tiles.
+
+## Files
+
+| file | content |
+|---|---|
+| `model.safetensors` | final weights after adversarial training (use this) |
+| `psrgan_x4_psnr.safetensors` | weights after the likelihood stage, smoother and higher PSNR |
+| `psrgan_x4_swag.safetensors` | SWAG mean and variance (32 snapshots) |
+| `config.json`, `psrgan.py` | model hyperparameters and code |
+| `eval/metrics.json` | every number on this page |
+| `notebooks/` | dataset creation and training notebooks |
+
+# Evaluation
+
+All numbers are on the test split of [PSRGAN Sentinel-2](https://huggingface.co/datasets/marklpd/PSRGAN-Sentinel-2) (1,056 patches), whose tiles lie in 0.5° blocks that share no pixel with training or validation. PSNR, SSIM and LPIPS are averaged per image (± is the 95% confidence interval), VGG is the perceptual loss used in training and OT the patchwise Sinkhorn distance between pixel color distributions.
+
+| method | PSNR (dB) ↑ | SSIM ↑ | LPIPS ↓ | VGG ↓ | OT ↓ |
+|---|---:|---:|---:|---:|---:|
+| Bicubic | 21.56 ± 0.17 | 0.5428 ± 0.0077 | 0.5987 ± 0.0050 | 8.135 | 0.00321 |
+| Real-ESRGAN | 20.22 ± 0.15 | 0.4753 ± 0.0068 | 0.3482 ± 0.0043 | 8.640 | 0.00399 |
+| PSRGAN (PSNR) | **21.94 ± 0.17** | **0.5981 ± 0.0063** | 0.2082 ± 0.0027 | **4.877** | 0.00267 |
+| **PSRGAN** | 21.28 ± 0.17 | 0.5694 ± 0.0063 | **0.1879 ± 0.0026** | 5.859 | **0.00135** |
+
+PSRGAN beats Real-ESRGAN in PSNR on 97% of test patches and in LPIPS on 100%. Compared with the likelihood stage, the adversarial stage gives up 0.66 dB PSNR and changes LPIPS from 0.208 to 0.188 (lower is better). Bicubic upsampling keeps a higher PSNR than the final model on 70% of patches, because PSNR rewards smooth averages over sharp but slightly misplaced detail. LPIPS, which follows perceived similarity, is lower for PSRGAN than for bicubic on 100% of patches.
+
+On the 128 test patches from Herzogenaurach, Wolfsburg and München:
+
+| method | PSNR (dB) ↑ | SSIM ↑ | LPIPS ↓ | VGG ↓ | OT ↓ |
+|---|---:|---:|---:|---:|---:|
+| Bicubic | 21.15 ± 0.25 | 0.5522 ± 0.0206 | 0.6825 ± 0.0142 | 9.868 | 0.00321 |
+| Real-ESRGAN | 20.24 ± 0.23 | 0.5147 ± 0.0190 | 0.3988 ± 0.0150 | 9.620 | 0.00308 |
+| PSRGAN (PSNR) | **21.30 ± 0.28** | **0.5960 ± 0.0199** | 0.2118 ± 0.0086 | **5.983** | 0.00310 |
+| **PSRGAN** | 20.61 ± 0.26 | 0.5620 ± 0.0188 | **0.1938 ± 0.0087** | 7.124 | **0.00146** |
+
+<center>
+  <img src="images/psrgan_city.png" width="900px">
+  <br><em>A test patch from Herzogenaurach.</em>
+</center>
+
+## Calibration
+
+| | ECE (σ vs. RMSE) | error within 1σ | error within 2σ |
+|---|---:|---:|---:|
+| Aleatoric + SWAG | 0.0240 | 86.6% | 98.4% |
+| Aleatoric + Bayesian heads | 0.0289 | 88.4% | 98.7% |
+
+A perfectly calibrated Gaussian would cover 68.3% and 95.4%. Both variants cover more, and the predicted σ is larger than the observed RMSE in every calibration bin, so the uncertainty is conservative. The uncertainty-error correlation is stronger in flat areas (Pearson r = 0.59) than on edges (r = 0.37).
+
+# Training
+
+PSRGAN starts from the Real-ESRGAN x4plus generator (RRDBNet, 23 blocks) and its U-Net discriminator.
+
+1. **Likelihood stage (768 epochs of 64 batches, batch size 16).** 0.6 · Gaussian NLL + 0.7 · (1 − SSIM) + 1.0 · VGG19 perceptual loss. All parameters loaded from Real-ESRGAN stay frozen for the first 168 epochs, so the new heads can settle first.
+2. **Adversarial stage (64 epochs).** Relativistic GAN loss (weight 0.01) plus contextual bilateral loss (weight 0.2), Adam with a learning rate of 2·10⁻⁴ halved every 32 epochs. SWAG collects a weight snapshot every 16 batches during the last 8 epochs.
+
+Augmentation (flips, 90° rotations) is applied to the training split only, and the test split is only used for the final evaluation. Training ran on one NVIDIA A100-SXM4-40GB in Google Colab.
+
+# Limitations
+
+- The model is trained on synthetically degraded Sentinel-2 images. It sharpens Sentinel-2 data at its own resolution and does not recover detail from a finer sensor. Details beyond what the input contains are estimates, not measurements.
+- Its training data are 98.6% Alpine tiles (Switzerland, Austria, northern Italy, south-eastern France, southern Germany, Slovenia) from late September and early October 2025. Expect worse results on lowland farmland, other climates, other seasons, deserts and dense city centres.
+- The uncertainty maps rank pixels by expected error, but the absolute σ is too large: 87% of pixel errors fall within 1σ instead of 68%. Treat σ as a conservative bound, or recalibrate it on your own data.
+- The final model trades pixel accuracy for sharpness. Its PSNR is 0.29 dB below bicubic upsampling, so use `psrgan_x4_psnr.safetensors` when pixel accuracy matters more than texture.
+- The training pixels are not square (about 10 m east-west, 15.6 m north-south), because tiles were requested on a square WGS84 grid.
+
+# Out-of-Scope Use
+
+- Measuring, mapping or deciding about individual parcels, buildings or people from super-resolved pixels.
+- Replacing actual higher-resolution imagery in legal, insurance, military or safety-critical decisions.
+- Presenting super-resolved images as original satellite observations.
+
+# License
+
+The code and weights are released under the BSD 3-Clause License, the same license as Real-ESRGAN, on which the weights are based. Training data contains modified Copernicus Sentinel data (2025).
+
+# Citation
+
+```
+@misc{leupold2026psrgan,
+  author = {Leupold, Mark-Daniel},
+  title = {PSRGAN: Probabilistic Super-Resolution for Sentinel-2 with Uncertainty Estimates},
+  year = {2026},
+  publisher = {Hugging Face},
+  url = {https://huggingface.co/marklpd/PSRGAN}
+}
+```

config.json

@@ -0,0 +1,9 @@
+{
+  "num_feat": 64,
+  "num_block": 23,
+  "num_grow_ch": 32,
+  "scale": 4,
+  "bayesian_heads": true,
+  "logvar_min": -10.0,
+  "logvar_max": 10.0
+}

eval/metrics.json

@@ -0,0 +1,247 @@
+{
+  "test": {
+    "n_patches": 1056,
+    "results": {
+      "Bicubic": {
+        "PSNR": 21.5607197952993,
+        "PSNR \u00b1": 0.1711068225145145,
+        "SSIM": 0.5427891769813318,
+        "SSIM \u00b1": 0.007653797755572996,
+        "LPIPS": 0.5987126496479367,
+        "LPIPS \u00b1": 0.004994828619425444,
+        "VGG": 8.134999072912967,
+        "OT": 0.003209914177811394
+      },
+      "Real-ESRGAN": {
+        "PSNR": 20.215020929322097,
+        "PSNR \u00b1": 0.15033991375476158,
+        "SSIM": 0.47526671260482434,
+        "SSIM \u00b1": 0.006763881811721902,
+        "LPIPS": 0.34822498551202996,
+        "LPIPS \u00b1": 0.004308759976577521,
+        "VGG": 8.63969765287457,
+        "OT": 0.00399009935614286
+      },
+      "PSRGAN (PSNR)": {
+        "PSNR": 21.936194533651527,
+        "PSNR \u00b1": 0.17383839192396056,
+        "SSIM": 0.5981000966646455,
+        "SSIM \u00b1": 0.006328675695824681,
+        "LPIPS": 0.20817970013895043,
+        "LPIPS \u00b1": 0.0027403351497943132,
+        "VGG": 4.8766192811908144,
+        "OT": 0.002672349598442176
+      },
+      "PSRGAN": {
+        "PSNR": 21.275709489981335,
+        "PSNR \u00b1": 0.16942986974481397,
+        "SSIM": 0.5693713306929126,
+        "SSIM \u00b1": 0.006270623736088413,
+        "LPIPS": 0.18790546653679374,
+        "LPIPS \u00b1": 0.002612060048786746,
+        "VGG": 5.858746337168144,
+        "OT": 0.0013482042025619496
+      }
+    },
+    "psrgan_win_rate": {
+      "Bicubic": {
+        "PSNR": 0.29829545454545453,
+        "SSIM": 0.6903409090909091,
+        "LPIPS": 1.0
+      },
+      "Real-ESRGAN": {
+        "PSNR": 0.9706439393939394,
+        "SSIM": 1.0,
+        "LPIPS": 1.0
+      },
+      "PSRGAN (PSNR)": {
+        "PSNR": 0.000946969696969697,
+        "SSIM": 0.002840909090909091,
+        "LPIPS": 0.9261363636363636
+      }
+    }
+  },
+  "test_city": {
+    "n_patches": 128,
+    "results": {
+      "Bicubic": {
+        "PSNR": 21.145912691950798,
+        "PSNR \u00b1": 0.2489663205630983,
+        "SSIM": 0.5521608623676002,
+        "SSIM \u00b1": 0.020581239817082063,
+        "LPIPS": 0.6825457969680429,
+        "LPIPS \u00b1": 0.014193458656650901,
+        "VGG": 9.868039011955261,
+        "OT": 0.003211126459063962
+      },
+      "Real-ESRGAN": {
+        "PSNR": 20.24102108180523,
+        "PSNR \u00b1": 0.22995863706651176,
+        "SSIM": 0.5146662218030542,
+        "SSIM \u00b1": 0.01896918156259728,
+        "LPIPS": 0.39876320981420577,
+        "LPIPS \u00b1": 0.014998464281860064,
+        "VGG": 9.620419383049011,
+        "OT": 0.003078749228734523
+      },
+      "PSRGAN (PSNR)": {
+        "PSNR": 21.299197003245354,
+        "PSNR \u00b1": 0.27709311095960687,
+        "SSIM": 0.5960347084328532,
+        "SSIM \u00b1": 0.019889227648886978,
+        "LPIPS": 0.21178198768757284,
+        "LPIPS \u00b1": 0.00864762965393086,
+        "VGG": 5.9829049706459045,
+        "OT": 0.003099578752880916
+      },
+      "PSRGAN": {
+        "PSNR": 20.610211864113808,
+        "PSNR \u00b1": 0.2630898367096838,
+        "SSIM": 0.5620126791764051,
+        "SSIM \u00b1": 0.018756582941474018,
+        "LPIPS": 0.19378441805019975,
+        "LPIPS \u00b1": 0.008681403045909196,
+        "VGG": 7.124138653278351,
+        "OT": 0.001463380001951009
+      }
+    }
+  },
+  "calibration": {
+    "swag": {
+      "ECE": 0.02397674263321663,
+      "coverage_1sigma": 0.8664602438608805,
+      "coverage_2sigma": 0.9841360251108805,
+      "reliability": {
+        "pred_std": [
+          0.020123129710555077,
+          0.042232051491737366,
+          0.06536165624856949,
+          0.08905356377363205,
+          0.11278203874826431,
+          0.1365380734205246,
+          0.16020965576171875,
+          0.18388722836971283,
+          0.2074737250804901,
+          0.2310594618320465,
+          0.2546865940093994,
+          0.2783450484275818,
+          0.3019865155220032,
+          0.32563266158103943,
+          0.3493105471134186
+        ],
+        "rmse": [
+          0.014095807448029518,
+          0.02998480387032032,
+          0.04914991185069084,
+          0.0686262995004654,
+          0.08680041134357452,
+          0.1045115739107132,
+          0.12249297648668289,
+          0.14000065624713898,
+          0.1562192142009735,
+          0.17163224518299103,
+          0.18531100451946259,
+          0.19668403267860413,
+          0.20926719903945923,
+          0.21913963556289673,
+          0.2296137809753418
+        ]
+      }
+    },
+    "bayesian_heads": {
+      "ECE": 0.02886166655259949,
+      "coverage_1sigma": 0.8842587471008301,
+      "coverage_2sigma": 0.986927350362142,
+      "reliability": {
+        "pred_std": [
+          0.02026936411857605,
+          0.04286534711718559,
+          0.0662725418806076,
+          0.08985785394906998,
+          0.11369309574365616,
+          0.1376895308494568,
+          0.16160856187343597,
+          0.18553175032138824,
+          0.20935258269309998,
+          0.233138307929039,
+          0.25700125098228455,
+          0.28088292479515076,
+          0.30468234419822693,
+          0.32858404517173767,
+          0.35245150327682495
+        ],
+        "rmse": [
+          0.012933187186717987,
+          0.027884116396307945,
+          0.045200906693935394,
+          0.0637541264295578,
+          0.08282162249088287,
+          0.10149621963500977,
+          0.1198882907629013,
+          0.13830198347568512,
+          0.15485383570194244,
+          0.17070721089839935,
+          0.1843496859073639,
+          0.19707782566547394,
+          0.2094695121049881,
+          0.21937879920005798,
+          0.23037448525428772
+        ]
+      }
+    }
+  },
+  "uncertainty_error_correlation": {
+    "pearson": 0.5869906544685364,
+    "spearman": 0.7000287473727114,
+    "by_region": {
+      "edges": {
+        "pearson": 0.3714364469051361,
+        "spearman": 0.37498411900628265
+      },
+      "flat": {
+        "pearson": 0.5925312638282776,
+        "spearman": 0.6675135785054067
+      }
+    }
+  },
+  "error_by_uncertainty_decile": [
+    0.009789826348423958,
+    0.016854647547006607,
+    0.024157894775271416,
+    0.03273877874016762,
+    0.04215969517827034,
+    0.052353471517562866,
+    0.06426148116588593,
+    0.07902297377586365,
+    0.09887537360191345,
+    0.13998572528362274
+  ],
+  "training": {
+    "train_patches": 8171,
+    "val_patches": 1008,
+    "test_patches": 1056,
+    "batch_size": 16,
+    "batches_per_epoch": 64,
+    "pretrain_epochs": 768,
+    "freeze_rrdb_epochs": 168,
+    "adversarial_epochs": 64,
+    "lr_G": 0.0002,
+    "lr_D": 0.0002,
+    "loss_weights_pretrain": {
+      "NLL": 0.6,
+      "SSIM": 0.7,
+      "VGG": 1.0
+    },
+    "loss_weights_adversarial": {
+      "ADV": 0.01,
+      "CX": 0.2
+    },
+    "swag": {
+      "snapshots": 32,
+      "start_epoch": 56,
+      "collect_every_batches": 16
+    },
+    "gpu": "NVIDIA A100-SXM4-40GB",
+    "seed": 42
+  }
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8708e583a031ef01cd9d3cb29a6bdce842f1ba1a64ac5ed9e9ee68e995c17f58
+size 66892468

notebooks/PSRGAN_dataset.ipynb

@@ -0,0 +1,1334 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# PSRGAN – Datensatz\n",
+    "\n",
+    "`Sentinel-2 Kacheln laden, Splits auf Kachelebene festlegen und LR/HR-Patches für Hugging Face erzeugen.`\n",
+    "\n",
+    "Ablauf: Kacheln (alt + neu) → Split pro Kachel (räumlich getrennt, 0,5°-Blöcke) → Patches → Parquet in `hf/dataset`.\n",
+    "Läuft auf CPU-Runtime, keine GPU nötig."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 0. Setup"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Installationen"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "!pip install rasterio cartopy --quiet"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import os\n",
+    "import io\n",
+    "import csv\n",
+    "import json\n",
+    "import math\n",
+    "import time\n",
+    "import random\n",
+    "import shutil\n",
+    "import threading\n",
+    "import warnings\n",
+    "import logging\n",
+    "from glob import glob\n",
+    "from io import BytesIO\n",
+    "from concurrent.futures import ThreadPoolExecutor\n",
+    "\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import cv2\n",
+    "import requests\n",
+    "import rasterio\n",
+    "from rasterio.env import Env\n",
+    "from PIL import Image\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import datasets\n",
+    "from datasets import Dataset, Features, Value, Array2D, Array3D\n",
+    "from datasets import Image as HFImage\n",
+    "\n",
+    "print(\"datasets:\", datasets.__version__)"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Google Drive Mounten"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from google.colab import drive\n",
+    "drive.mount(\"/content/drive\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "logging.getLogger(\"rasterio\").setLevel(logging.ERROR)\n",
+    "logging.getLogger(\"rasterio._env\").setLevel(logging.ERROR)\n",
+    "\n",
+    "gdal_env = Env(GDAL_DISABLE_READDIR_ON_OPEN=\"YES\")\n",
+    "gdal_env.__enter__()\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\", category=rasterio.errors.NotGeoreferencedWarning)\n",
+    "warnings.filterwarnings(\"ignore\", message=\".*TIFFReadDirectory.*ExtraSamples.*\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 1. Konfiguration"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "config = {\n",
+    "    # Sentinel-2\n",
+    "    \"BANDS\": [\"B4\", \"B3\", \"B2\"],\n",
+    "\n",
+    "    # Geometrie\n",
+    "    \"PATCH_SIZE\": 256,\n",
+    "    \"UPSCALE\": 4,\n",
+    "    \"LR_PATCH_SIZE\": 64,\n",
+    "\n",
+    "    # AOI\n",
+    "    \"AOI_KM\": 8.0,\n",
+    "    \"REQUEST_PIXELS\": 512,\n",
+    "\n",
+    "    # Sampling\n",
+    "    \"TARGET_PATCHES\": 10000,\n",
+    "    \"PATCHES_PER_TILE\": 16,\n",
+    "    \"MAX_CLOUD_TILE\": 0.25,\n",
+    "    \"MAX_CLOUD_PATCH\": 0.15,\n",
+    "    \"MAX_WATER_TILE\": 0.5,    # neu: Meer/Seen raus\n",
+    "    \"MAX_NODATA_TILE\": 0.01,  # neu: Swath-Rand raus\n",
+    "\n",
+    "    # Region (alt, erste 192 Kacheln)\n",
+    "    \"REGION_V1\": {\n",
+    "        \"lon_min\": 7.5,\n",
+    "        \"lon_max\": 13.5,\n",
+    "        \"lat_min\": 46.5,\n",
+    "        \"lat_max\": 51.5,\n",
+    "    },\n",
+    "\n",
+    "    # Region (neu, Mitteleuropa)\n",
+    "    \"REGION\": {\n",
+    "        \"lon_min\": 4.0,\n",
+    "        \"lon_max\": 19.5,\n",
+    "        \"lat_min\": 45.5,\n",
+    "        \"lat_max\": 54.5,\n",
+    "    },\n",
+    "\n",
+    "    # Zeitfenster (wie bisher)\n",
+    "    \"TIME_FROM\": \"2025-09-20T00:00:00Z\",\n",
+    "    \"TIME_TO\": \"2025-10-07T23:59:59Z\",\n",
+    "\n",
+    "    # Splits\n",
+    "    \"SPLIT_FRACTIONS\": {\"train\": 0.8, \"validation\": 0.1, \"test\": 0.1},\n",
+    "    \"BLOCK_DEG\": 0.5,\n",
+    "\n",
+    "    # Download\n",
+    "    \"N_WORKERS\": 6,\n",
+    "\n",
+    "    # Zufall\n",
+    "    \"SEED\": 42,\n",
+    "}\n",
+    "\n",
+    "config[\"AOI_SIZE_DEG\"] = config[\"AOI_KM\"] / 111.0\n",
+    "config[\"LR_PATCH_SIZE\"] = config[\"PATCH_SIZE\"] // config[\"UPSCALE\"]\n",
+    "\n",
+    "for k, v in config.items():\n",
+    "    print(f\"{k:22s}: {v}\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "def set_seed(seed: int):\n",
+    "    random.seed(seed)\n",
+    "    np.random.seed(seed)\n",
+    "    print(f\"Seed set to {seed}\")\n",
+    "\n",
+    "set_seed(config[\"SEED\"])"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Pfade"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "DATA_ROOT = \"/content/drive/MyDrive/OT_SRGAN/data\"\n",
+    "DOWNLOAD_DIR = os.path.join(DATA_ROOT, \"s2_aoi_tiles\")\n",
+    "CUSTOM_DOWNLOAD_DIR = os.path.join(DATA_ROOT, \"custom_s2_tiles\")\n",
+    "SPLITS_PATH = os.path.join(DATA_ROOT, \"tile_splits.csv\")\n",
+    "\n",
+    "HF_DATASET_DIR = \"/content/drive/MyDrive/OT_SRGAN/hf/dataset\"\n",
+    "HF_DATA_DIR = os.path.join(HF_DATASET_DIR, \"data\")\n",
+    "HF_IMAGES_DIR = os.path.join(HF_DATASET_DIR, \"images\")\n",
+    "HF_META_DIR = os.path.join(HF_DATASET_DIR, \"metadata\")\n",
+    "\n",
+    "# Lokal bauen, dann auf Drive kopieren\n",
+    "LOCAL_OUT = \"/content/hf_dataset_build\"\n",
+    "LOCAL_CACHE = \"/content/hf_cache\"\n",
+    "\n",
+    "for d in [DOWNLOAD_DIR, CUSTOM_DOWNLOAD_DIR, HF_DATA_DIR, HF_IMAGES_DIR, HF_META_DIR, LOCAL_OUT, LOCAL_CACHE]:\n",
+    "    os.makedirs(d, exist_ok=True)\n",
+    "\n",
+    "PROCESS_URL = \"https://sh.dataspace.copernicus.eu/api/v1/process\""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 2. Kacheln laden"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Token"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from google.colab import userdata\n",
+    "\n",
+    "CLIENT_ID = userdata.get(\"COPERNICUS_CLIENT_ID\")\n",
+    "CLIENT_SECRET = userdata.get(\"COPERNICUS_CLIENT_SECRET\")\n",
+    "\n",
+    "token_url = (\n",
+    "    \"https://identity.dataspace.copernicus.eu/auth/\"\n",
+    "    \"realms/CDSE/protocol/openid-connect/token\"\n",
+    ")\n",
+    "\n",
+    "HEADERS = {\"Content-Type\": \"application/json\"}\n",
+    "token_lock = threading.Lock()\n",
+    "\n",
+    "def refresh_token():\n",
+    "    # Token läuft nach ~10 min ab\n",
+    "    resp = requests.post(\n",
+    "        token_url,\n",
+    "        data={\n",
+    "            \"grant_type\": \"client_credentials\",\n",
+    "            \"client_id\": CLIENT_ID,\n",
+    "            \"client_secret\": CLIENT_SECRET,\n",
+    "        },\n",
+    "    )\n",
+    "    resp.raise_for_status()\n",
+    "    HEADERS[\"Authorization\"] = f\"Bearer {resp.json()['access_token']}\"\n",
+    "\n",
+    "refresh_token()\n",
+    "print(\"Access token acquired.\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Request"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "EVALSCRIPT = \"\"\"\n",
+    "//VERSION=3\n",
+    "function setup() {\n",
+    "  return {\n",
+    "    input: [{bands:[\"B02\",\"B03\",\"B04\",\"B08\",\"SCL\"]}],\n",
+    "    output: {bands: 5, sampleType: \"FLOAT32\"}\n",
+    "  };\n",
+    "}\n",
+    "\n",
+    "function evaluatePixel(s) {\n",
+    "  let ndvi = (s.B08 - s.B04) / (s.B08 + s.B04 + 1e-6);\n",
+    "  return [s.B04, s.B03, s.B02, ndvi, s.SCL];\n",
+    "}\n",
+    "\"\"\""
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "def build_request(bbox):\n",
+    "    return {\n",
+    "        \"input\": {\n",
+    "            \"bounds\": {\"bbox\": bbox},\n",
+    "            \"data\": [{\n",
+    "                \"type\": \"sentinel-2-l2a\",\n",
+    "                \"dataFilter\": {\n",
+    "                    \"timeRange\": {\n",
+    "                        \"from\": config[\"TIME_FROM\"],\n",
+    "                        \"to\":   config[\"TIME_TO\"],\n",
+    "                    }\n",
+    "                },\n",
+    "                \"processing\": {\"mosaickingOrder\": \"leastCC\"},\n",
+    "            }],\n",
+    "        },\n",
+    "        \"output\": {\n",
+    "            \"width\": config[\"REQUEST_PIXELS\"],\n",
+    "            \"height\": config[\"REQUEST_PIXELS\"],\n",
+    "            \"responses\": [{\n",
+    "                \"identifier\": \"default\",\n",
+    "                \"format\": {\"type\": \"image/tiff\"},\n",
+    "            }],\n",
+    "        },\n",
+    "        \"evalscript\": EVALSCRIPT,\n",
+    "    }"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "def compute_cloud_mask(scl):\n",
+    "    cloud_classes = {3,7,8,9,10}\n",
+    "    return np.isin(scl, list(cloud_classes))\n",
+    "\n",
+    "def compute_cloud_ratio(scl):\n",
+    "    mask = compute_cloud_mask(scl)\n",
+    "    return mask.sum() / scl.size\n",
+    "\n",
+    "def parse_multiband_tiff(content):\n",
+    "    with rasterio.open(BytesIO(content)) as src:\n",
+    "        rgb = src.read([1,2,3]).astype(np.float32)\n",
+    "        ndvi = src.read(4).astype(np.float32)\n",
+    "        scl = src.read(5).astype(np.float32)\n",
+    "    return rgb, ndvi, scl\n",
+    "\n",
+    "def bbox_overlap(a, b):\n",
+    "    # [lon_min, lat_min, lon_max, lat_max]\n",
+    "    return not (a[2] <= b[0] or b[2] <= a[0] or a[3] <= b[1] or b[3] <= a[1])"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Bestehende Kacheln"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "existing_tiles = sorted(glob(os.path.join(DOWNLOAD_DIR, \"aoi_*.npz\")))\n",
+    "custom_tiles = sorted(glob(os.path.join(CUSTOM_DOWNLOAD_DIR, \"custom_*.npz\")))\n",
+    "\n",
+    "occupied = [list(np.load(p)[\"bbox\"]) for p in tqdm(existing_tiles + custom_tiles, desc=\"Read bboxes\", ncols=120)]\n",
+    "\n",
+    "n_tiles_target = math.ceil(config[\"TARGET_PATCHES\"] / config[\"PATCHES_PER_TILE\"] * 1.03)  # 3% Puffer für verworfene Patches\n",
+    "n_new = max(0, n_tiles_target - len(existing_tiles) - len(custom_tiles))\n",
+    "\n",
+    "print(f\"Found {len(existing_tiles)} AOI tiles + {len(custom_tiles)} custom tiles.\")\n",
+    "print(f\"Target {n_tiles_target} tiles → {n_new} new tiles to download.\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Neue Kacheln (parallel)\n",
+    "\n",
+    "Neue Kacheln überlappen keine bestehende Kachel. Bei erneutem Ausführen wird nur der Rest geladen."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "state_lock = threading.Lock()\n",
+    "rng_bbox = random.Random(config[\"SEED\"] + len(existing_tiles))\n",
+    "next_idx = max([int(os.path.basename(p)[4:8]) for p in existing_tiles], default=-1) + 1\n",
+    "collected = []\n",
+    "rejected = {\"http\": 0, \"cloud\": 0, \"water\": 0, \"nodata\": 0, \"parse\": 0}\n",
+    "\n",
+    "def random_bbox():\n",
+    "    r = config[\"REGION\"]\n",
+    "    lon = rng_bbox.uniform(r[\"lon_min\"], r[\"lon_max\"])\n",
+    "    lat = rng_bbox.uniform(r[\"lat_min\"], r[\"lat_max\"])\n",
+    "    d = config[\"AOI_SIZE_DEG\"] / 2\n",
+    "    return [lon - d, lat - d, lon + d, lat + d]\n",
+    "\n",
+    "def reserve_bbox():\n",
+    "    # Freie Fläche suchen und reservieren\n",
+    "    with state_lock:\n",
+    "        while True:\n",
+    "            bbox = random_bbox()\n",
+    "            if not any(bbox_overlap(bbox, o) for o in occupied):\n",
+    "                occupied.append(bbox)\n",
+    "                return bbox\n",
+    "\n",
+    "def release_bbox(bbox):\n",
+    "    with state_lock:\n",
+    "        occupied.remove(bbox)\n",
+    "\n",
+    "def post_with_retry(body, max_tries=6):\n",
+    "    for t in range(max_tries):\n",
+    "        r = requests.post(PROCESS_URL, headers=dict(HEADERS), json=body, timeout=120)\n",
+    "        if r.status_code == 401:\n",
+    "            with token_lock:\n",
+    "                refresh_token()\n",
+    "            continue\n",
+    "        if r.status_code == 429 or r.status_code >= 500:\n",
+    "            time.sleep(2 ** t)\n",
+    "            continue\n",
+    "        return r\n",
+    "    return r\n",
+    "\n",
+    "def download_worker(pbar):\n",
+    "    global next_idx\n",
+    "    while True:\n",
+    "        with state_lock:\n",
+    "            if len(collected) >= n_new:\n",
+    "                return\n",
+    "\n",
+    "        bbox = reserve_bbox()\n",
+    "        try:\n",
+    "            r = post_with_retry(build_request(bbox))\n",
+    "        except requests.RequestException:\n",
+    "            r = None\n",
+    "        if r is None or r.status_code != 200:\n",
+    "            rejected[\"http\"] += 1\n",
+    "            release_bbox(bbox)\n",
+    "            if r is not None and r.status_code in (402, 403):\n",
+    "                print(\"Quota/permission error:\", r.status_code, r.text[:200])\n",
+    "                return\n",
+    "            continue\n",
+    "\n",
+    "        try:\n",
+    "            rgb, ndvi, scl = parse_multiband_tiff(r.content)\n",
+    "        except Exception:\n",
+    "            rejected[\"parse\"] += 1\n",
+    "            release_bbox(bbox)\n",
+    "            continue\n",
+    "\n",
+    "        if compute_cloud_ratio(scl) > config[\"MAX_CLOUD_TILE\"]:\n",
+    "            rejected[\"cloud\"] += 1; release_bbox(bbox); continue\n",
+    "        if (scl == 6).mean() > config[\"MAX_WATER_TILE\"]:\n",
+    "            rejected[\"water\"] += 1; release_bbox(bbox); continue\n",
+    "        if np.isin(scl, [0, 1]).mean() > config[\"MAX_NODATA_TILE\"]:\n",
+    "            rejected[\"nodata\"] += 1; release_bbox(bbox); continue\n",
+    "\n",
+    "        with state_lock:\n",
+    "            if len(collected) >= n_new:\n",
+    "                release_bbox_unlocked = True\n",
+    "            else:\n",
+    "                release_bbox_unlocked = False\n",
+    "                tile_id = f\"aoi_{next_idx:04d}\"\n",
+    "                next_idx += 1\n",
+    "                collected.append(tile_id)\n",
+    "        if release_bbox_unlocked:\n",
+    "            release_bbox(bbox)\n",
+    "            return\n",
+    "\n",
+    "        np.savez_compressed(\n",
+    "            os.path.join(DOWNLOAD_DIR, f\"{tile_id}.npz\"),\n",
+    "            rgb=rgb,\n",
+    "            ndvi=ndvi,\n",
+    "            scl=scl,\n",
+    "            bbox=bbox,\n",
+    "            time_from=config[\"TIME_FROM\"],\n",
+    "            time_to=config[\"TIME_TO\"],\n",
+    "        )\n",
+    "        pbar.update(1)\n",
+    "\n",
+    "if n_new > 0:\n",
+    "    pbar = tqdm(total=n_new, desc=\"Downloading AOIs\", ncols=120)\n",
+    "    with ThreadPoolExecutor(max_workers=config[\"N_WORKERS\"]) as ex:\n",
+    "        for _ in range(config[\"N_WORKERS\"]):\n",
+    "            ex.submit(download_worker, pbar)\n",
+    "    pbar.close()\n",
+    "\n",
+    "print(f\"Added {len(collected)} new AOI tiles. Rejected: {rejected}\")\n",
+    "\n",
+    "existing_tiles = sorted(glob(os.path.join(DOWNLOAD_DIR, \"aoi_*.npz\")))\n",
+    "print(f\"Total AOI tiles on disk: {len(existing_tiles)}\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 3. Splits auf Kachelebene\n",
+    "\n",
+    "1. Überlappende Kacheln → eine Gruppe (Union-Find).\n",
+    "2. Gruppen nach 0,5°-Blöcken zusammenfassen, Blöcke zufällig auf Splits verteilen (räumlich getrennt).\n",
+    "3. Blöcke mit Stadt-Kacheln (Herzogenaurach, Wolfsburg, München) → immer `test`.\n",
+    "\n",
+    "Augmentation passiert erst im Training und nur auf `train`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# [lon_min, lat_min, lon_max, lat_max]\n",
+    "CUSTOM_AOIS = {\n",
+    "    \"Herzogenaurach\": [\n",
+    "        [10.85, 49.55, 10.95, 49.65],\n",
+    "        [10.80, 49.57, 10.90, 49.67],\n",
+    "        [10.87, 49.53, 10.97, 49.63],\n",
+    "        [10.83, 49.56, 10.93, 49.66],\n",
+    "    ],\n",
+    "    \"Wolfsburg\": [\n",
+    "        [10.75, 52.40, 10.85, 52.50],\n",
+    "        [10.70, 52.42, 10.80, 52.52],\n",
+    "    ],\n",
+    "    \"München\": [\n",
+    "        [11.50, 48.10, 11.60, 48.20],\n",
+    "        [11.55, 48.12, 11.65, 48.22],\n",
+    "    ],\n",
+    "}\n",
+    "\n",
+    "def city_name(bbox):\n",
+    "    for name, boxes in CUSTOM_AOIS.items():\n",
+    "        if any(np.allclose(bbox, b) for b in boxes):\n",
+    "            return name\n",
+    "    return \"unknown\"\n",
+    "\n",
+    "tiles = []\n",
+    "for p in existing_tiles:\n",
+    "    tiles.append({\"tile_id\": os.path.basename(p)[:-4], \"path\": p, \"source\": \"random\", \"aoi_name\": \"\",\n",
+    "                  \"bbox\": [float(v) for v in np.load(p)[\"bbox\"]]})\n",
+    "for p in custom_tiles:\n",
+    "    bbox = [float(v) for v in np.load(p)[\"bbox\"]]\n",
+    "    tiles.append({\"tile_id\": os.path.basename(p)[:-4], \"path\": p, \"source\": \"city\", \"aoi_name\": city_name(bbox),\n",
+    "                  \"bbox\": bbox})\n",
+    "\n",
+    "print(f\"{len(tiles)} tiles total\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Union-Find über überlappende Kacheln\n",
+    "parent = list(range(len(tiles)))\n",
+    "\n",
+    "def find(i):\n",
+    "    while parent[i] != i:\n",
+    "        parent[i] = parent[parent[i]]\n",
+    "        i = parent[i]\n",
+    "    return i\n",
+    "\n",
+    "for i in range(len(tiles)):\n",
+    "    for j in range(i + 1, len(tiles)):\n",
+    "        if bbox_overlap(tiles[i][\"bbox\"], tiles[j][\"bbox\"]):\n",
+    "            parent[find(i)] = find(j)\n",
+    "\n",
+    "groups = {}\n",
+    "for i in range(len(tiles)):\n",
+    "    groups.setdefault(find(i), []).append(i)\n",
+    "\n",
+    "sizes = sorted([len(g) for g in groups.values()], reverse=True)\n",
+    "print(f\"{len(groups)} overlap groups, largest: {sizes[:5]}\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Blöcke\n",
+    "B = config[\"BLOCK_DEG\"]\n",
+    "\n",
+    "def block_of(idxs):\n",
+    "    lon = np.mean([(tiles[i][\"bbox\"][0] + tiles[i][\"bbox\"][2]) / 2 for i in idxs])\n",
+    "    lat = np.mean([(tiles[i][\"bbox\"][1] + tiles[i][\"bbox\"][3]) / 2 for i in idxs])\n",
+    "    return (int(math.floor(lon / B)), int(math.floor(lat / B)))\n",
+    "\n",
+    "block_groups = {}\n",
+    "for g, idxs in groups.items():\n",
+    "    block_groups.setdefault(block_of(idxs), []).append(g)\n",
+    "\n",
+    "block_tiles = {b: sum(len(groups[g]) for g in gs) for b, gs in block_groups.items()}\n",
+    "city_blocks = {b for b, gs in block_groups.items()\n",
+    "               if any(tiles[i][\"source\"] == \"city\" for g in gs for i in groups[g])}\n",
+    "\n",
+    "# Greedy: Block in den Split mit dem größten relativen Defizit\n",
+    "fractions = config[\"SPLIT_FRACTIONS\"]\n",
+    "target = {s: f * len(tiles) for s, f in fractions.items()}\n",
+    "count = {s: 0 for s in fractions}\n",
+    "block_split = {}\n",
+    "\n",
+    "for b in sorted(city_blocks):\n",
+    "    block_split[b] = \"test\"\n",
+    "    count[\"test\"] += block_tiles[b]\n",
+    "\n",
+    "rest = sorted(set(block_groups) - city_blocks)\n",
+    "np.random.default_rng(config[\"SEED\"]).shuffle(rest)\n",
+    "for b in rest:\n",
+    "    s = max(fractions, key=lambda s: (target[s] - count[s]) / target[s])\n",
+    "    block_split[b] = s\n",
+    "    count[s] += block_tiles[b]\n",
+    "\n",
+    "for g, idxs in groups.items():\n",
+    "    b = block_of(idxs)\n",
+    "    for i in idxs:\n",
+    "        tiles[i][\"split\"] = block_split[b]\n",
+    "        tiles[i][\"block\"] = f\"{b[0]}_{b[1]}\"\n",
+    "        tiles[i][\"group\"] = int(g)\n",
+    "\n",
+    "print(f\"{len(block_groups)} blocks ({len(city_blocks)} city blocks → test)\")\n",
+    "for s in fractions:\n",
+    "    print(f\"{s:10s}: {count[s]:4d} tiles ({count[s] / len(tiles):.1%})\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Check: keine Überlappung zwischen Splits\n",
+    "n_cross = 0\n",
+    "for i in range(len(tiles)):\n",
+    "    for j in range(i + 1, len(tiles)):\n",
+    "        if tiles[i][\"split\"] != tiles[j][\"split\"] and bbox_overlap(tiles[i][\"bbox\"], tiles[j][\"bbox\"]):\n",
+    "            n_cross += 1\n",
+    "assert n_cross == 0, f\"{n_cross} overlapping tile pairs across splits\"\n",
+    "assert all(t[\"split\"] == \"test\" for t in tiles if t[\"source\"] == \"city\")\n",
+    "print(\"OK: no pixel is shared between train, validation and test.\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "tiles_df = pd.DataFrame([{\n",
+    "    \"tile_id\": t[\"tile_id\"], \"split\": t[\"split\"], \"source\": t[\"source\"], \"aoi_name\": t[\"aoi_name\"],\n",
+    "    \"lon_min\": t[\"bbox\"][0], \"lat_min\": t[\"bbox\"][1], \"lon_max\": t[\"bbox\"][2], \"lat_max\": t[\"bbox\"][3],\n",
+    "    \"block\": t[\"block\"], \"group\": t[\"group\"],\n",
+    "} for t in tiles])\n",
+    "tiles_df.to_csv(SPLITS_PATH, index=False)\n",
+    "print(f\"Saved {SPLITS_PATH}\")\n",
+    "tiles_df.groupby([\"split\", \"source\"]).size()"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 4. Patches"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "PATCH_SIZE_HR = config[\"PATCH_SIZE\"]\n",
+    "UPSCALE = config[\"UPSCALE\"]\n",
+    "PATCH_SIZE_LR = config[\"LR_PATCH_SIZE\"]"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "def normalize_for_display(patch):\n",
+    "    p_low = np.percentile(patch, 1)\n",
+    "    p_high = np.percentile(patch, 99)\n",
+    "    norm = (patch - p_low) / (p_high - p_low + 1e-8)\n",
+    "    return np.clip(norm, 0, 1)"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "INTERP_NAMES = {cv2.INTER_CUBIC: \"bicubic\", cv2.INTER_LINEAR: \"bilinear\", cv2.INTER_AREA: \"area\"}\n",
+    "\n",
+    "def downsample_realistic(hr, scale, return_params=False):\n",
+    "    # HWC für OpenCV\n",
+    "    img = hr.transpose(1, 2, 0)\n",
+    "\n",
+    "    h, w = img.shape[:2]\n",
+    "\n",
+    "    # Zufallsfilter\n",
+    "    filter_type = random.choice([\"gaussian\", \"motion\", \"defocus\"])\n",
+    "    sigma = 0.0\n",
+    "\n",
+    "    if filter_type == \"gaussian\":\n",
+    "        # Isotropischer Gauss blur mit zufall kernel + sigma (SIGMA)\n",
+    "        k = random.choice([3, 5, 7])\n",
+    "        sigma = random.uniform(0.5, 1.5)\n",
+    "        img_blur = cv2.GaussianBlur(img, (k, k), sigmaX=sigma)\n",
+    "\n",
+    "    elif filter_type == \"motion\":\n",
+    "        # Anisotropischer motion blur kernel\n",
+    "        k = random.choice([5, 7, 9])\n",
+    "        kernel = np.zeros((k, k), dtype=np.float32)\n",
+    "\n",
+    "        # Zufall horizontal oder vertikal\n",
+    "        if random.random() < 0.5:\n",
+    "            kernel[k // 2, :] = 1.0\n",
+    "            filter_type = \"motion_h\"\n",
+    "        else:\n",
+    "            kernel[:, k // 2] = 1.0\n",
+    "            filter_type = \"motion_v\"\n",
+    "\n",
+    "        kernel /= k\n",
+    "        img_blur = cv2.filter2D(img, -1, kernel)\n",
+    "\n",
+    "    elif filter_type == \"defocus\":\n",
+    "        # Approximierter defocus blur mit kreiskernel\n",
+    "        k = random.choice([5, 7, 9])\n",
+    "        kernel = np.zeros((k, k), np.float32)\n",
+    "        center = k // 2\n",
+    "\n",
+    "        for i in range(k):\n",
+    "            for j in range(k):\n",
+    "                if (i - center) ** 2 + (j - center) ** 2 <= center ** 2:\n",
+    "                    kernel[i, j] = 1\n",
+    "\n",
+    "        kernel /= np.sum(kernel)\n",
+    "        img_blur = cv2.filter2D(img, -1, kernel)\n",
+    "\n",
+    "    # Zufallsinterpolation\n",
+    "    interpolation = random.choice([\n",
+    "        cv2.INTER_CUBIC,    # bikubisch\n",
+    "        cv2.INTER_LINEAR,   # bilinear\n",
+    "        cv2.INTER_AREA      # flächenbasiert\n",
+    "    ])\n",
+    "\n",
+    "    lr = cv2.resize(\n",
+    "        img_blur,\n",
+    "        (w // scale, h // scale),\n",
+    "        interpolation=interpolation\n",
+    "    )\n",
+    "\n",
+    "    # Zurück zu CHW\n",
+    "    lr = lr.transpose(2, 0, 1).astype(np.float32)\n",
+    "    if not return_params:\n",
+    "        return lr\n",
+    "    return lr, {\"blur\": filter_type, \"blur_kernel\": int(k), \"blur_sigma\": float(sigma),\n",
+    "                \"interpolation\": INTERP_NAMES[interpolation]}"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "def extract_patches_random(aoi_data, n_patches):\n",
+    "    rgb, ndvi, scl = aoi_data\n",
+    "    H, W = rgb.shape[1:]\n",
+    "\n",
+    "    extracted = []\n",
+    "    attempts = 0\n",
+    "    max_attempts = n_patches * 10\n",
+    "\n",
+    "    while len(extracted) < n_patches and attempts < max_attempts:\n",
+    "        y = np.random.randint(0, H - PATCH_SIZE_HR)\n",
+    "        x = np.random.randint(0, W - PATCH_SIZE_HR)\n",
+    "\n",
+    "        patch_rgb = rgb[:, y:y+PATCH_SIZE_HR, x:x+PATCH_SIZE_HR]\n",
+    "        patch_ndvi = ndvi[y:y+PATCH_SIZE_HR, x:x+PATCH_SIZE_HR]\n",
+    "        patch_scl = scl[y:y+PATCH_SIZE_HR, x:x+PATCH_SIZE_HR]\n",
+    "\n",
+    "        cloud_ratio = np.isin(\n",
+    "            patch_scl, [3,7,8,9,10]\n",
+    "        ).mean()\n",
+    "\n",
+    "        if cloud_ratio > config[\"MAX_CLOUD_PATCH\"]:\n",
+    "            attempts += 1\n",
+    "            continue\n",
+    "\n",
+    "        hr = normalize_for_display(patch_rgb / 10000.0)\n",
+    "        lr, deg = downsample_realistic(hr, UPSCALE, return_params=True)\n",
+    "\n",
+    "        extracted.append({\n",
+    "            \"lr\": lr, \"hr\": hr, \"ndvi\": patch_ndvi, \"scl\": patch_scl, \"rgb\": patch_rgb,\n",
+    "            \"y\": int(y), \"x\": int(x), \"cloud_ratio\": float(cloud_ratio),\n",
+    "            \"p_low\": float(np.percentile(patch_rgb, 1)), \"p_high\": float(np.percentile(patch_rgb, 99)),\n",
+    "            **deg,\n",
+    "        })\n",
+    "        attempts += 1\n",
+    "\n",
+    "    return extracted"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Beispiel"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "data = np.load(tiles[0][\"path\"])\n",
+    "np.random.seed(0); random.seed(0)\n",
+    "ex = extract_patches_random((data[\"rgb\"], data[\"ndvi\"], data[\"scl\"]), 4)\n",
+    "\n",
+    "plt.figure(figsize=(12, 9))\n",
+    "for i, p in enumerate(ex):\n",
+    "    plt.subplot(3, 4, i + 1); plt.imshow(p[\"hr\"].transpose(1, 2, 0)); plt.title(\"HR\"); plt.axis(\"off\")\n",
+    "    plt.subplot(3, 4, i + 5); plt.imshow(np.clip(p[\"lr\"], 0, 1).transpose(1, 2, 0)); plt.title(f\"LR ({p['blur']}, {p['interpolation']})\"); plt.axis(\"off\")\n",
+    "    plt.subplot(3, 4, i + 9); plt.imshow(p[\"ndvi\"], cmap=\"RdYlGn\", vmin=-1, vmax=1); plt.title(\"NDVI\"); plt.axis(\"off\")\n",
+    "plt.tight_layout(); plt.show()"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Hugging Face Schema"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "P = PATCH_SIZE_HR\n",
+    "\n",
+    "features = Features({\n",
+    "    \"patch_id\": Value(\"string\"),\n",
+    "    \"tile_id\": Value(\"string\"),\n",
+    "    \"source\": Value(\"string\"),\n",
+    "    \"aoi_name\": Value(\"string\"),\n",
+    "    \"lr\": HFImage(),\n",
+    "    \"hr\": HFImage(),\n",
+    "    \"ndvi\": Array2D((P, P), \"float32\"),\n",
+    "    \"scl\": Array2D((P, P), \"uint8\"),\n",
+    "    \"reflectance\": Array3D((3, P, P), \"uint16\"),\n",
+    "    \"lon\": Value(\"float64\"),\n",
+    "    \"lat\": Value(\"float64\"),\n",
+    "    \"lon_min\": Value(\"float64\"),\n",
+    "    \"lat_min\": Value(\"float64\"),\n",
+    "    \"lon_max\": Value(\"float64\"),\n",
+    "    \"lat_max\": Value(\"float64\"),\n",
+    "    \"crop_y\": Value(\"int32\"),\n",
+    "    \"crop_x\": Value(\"int32\"),\n",
+    "    \"cloud_fraction\": Value(\"float32\"),\n",
+    "    \"water_fraction\": Value(\"float32\"),\n",
+    "    \"stretch_low\": Value(\"float32\"),\n",
+    "    \"stretch_high\": Value(\"float32\"),\n",
+    "    \"blur\": Value(\"string\"),\n",
+    "    \"blur_kernel\": Value(\"int32\"),\n",
+    "    \"blur_sigma\": Value(\"float32\"),\n",
+    "    \"interpolation\": Value(\"string\"),\n",
+    "    \"time_from\": Value(\"string\"),\n",
+    "    \"time_to\": Value(\"string\"),\n",
+    "})\n",
+    "\n",
+    "def to_png(arr_chw):\n",
+    "    return Image.fromarray(np.round(np.clip(arr_chw, 0, 1).transpose(1, 2, 0) * 255).astype(np.uint8))\n",
+    "\n",
+    "def tile_rows(t, tile_index):\n",
+    "    data = np.load(t[\"path\"])\n",
+    "    rgb, ndvi, scl = data[\"rgb\"], data[\"ndvi\"], data[\"scl\"]\n",
+    "\n",
+    "    # Reproduzierbar pro Kachel\n",
+    "    seed = config[\"SEED\"] * 100003 + tile_index\n",
+    "    np.random.seed(seed); random.seed(seed)\n",
+    "    patches = extract_patches_random((rgb, ndvi, scl), n_patches=config[\"PATCHES_PER_TILE\"])\n",
+    "\n",
+    "    lon_min, lat_min, lon_max, lat_max = t[\"bbox\"]\n",
+    "    H, W = rgb.shape[1:]\n",
+    "    dlon, dlat = (lon_max - lon_min) / W, (lat_max - lat_min) / H\n",
+    "    t_from = str(data[\"time_from\"]) if \"time_from\" in data.files else (config[\"TIME_FROM\"] if t[\"source\"] == \"random\" else \"\")\n",
+    "    t_to = str(data[\"time_to\"]) if \"time_to\" in data.files else (config[\"TIME_TO\"] if t[\"source\"] == \"random\" else \"\")\n",
+    "\n",
+    "    for k, p in enumerate(patches):\n",
+    "        # Zeile 0 = Norden\n",
+    "        p_lon_min = lon_min + p[\"x\"] * dlon\n",
+    "        p_lon_max = p_lon_min + P * dlon\n",
+    "        p_lat_max = lat_max - p[\"y\"] * dlat\n",
+    "        p_lat_min = p_lat_max - P * dlat\n",
+    "        yield {\n",
+    "            \"patch_id\": f\"{t['tile_id']}_p{k:02d}\",\n",
+    "            \"tile_id\": t[\"tile_id\"],\n",
+    "            \"source\": t[\"source\"],\n",
+    "            \"aoi_name\": t[\"aoi_name\"],\n",
+    "            \"lr\": to_png(p[\"lr\"]),\n",
+    "            \"hr\": to_png(p[\"hr\"]),\n",
+    "            \"ndvi\": np.clip(p[\"ndvi\"], -1, 1).astype(np.float32),\n",
+    "            \"scl\": p[\"scl\"].astype(np.uint8),\n",
+    "            \"reflectance\": np.round(p[\"rgb\"] * 10000).clip(0, 65535).astype(np.uint16),\n",
+    "            \"lon\": (p_lon_min + p_lon_max) / 2,\n",
+    "            \"lat\": (p_lat_min + p_lat_max) / 2,\n",
+    "            \"lon_min\": p_lon_min, \"lat_min\": p_lat_min, \"lon_max\": p_lon_max, \"lat_max\": p_lat_max,\n",
+    "            \"crop_y\": p[\"y\"], \"crop_x\": p[\"x\"],\n",
+    "            \"cloud_fraction\": p[\"cloud_ratio\"],\n",
+    "            \"water_fraction\": float((p[\"scl\"] == 6).mean()),\n",
+    "            \"stretch_low\": p[\"p_low\"], \"stretch_high\": p[\"p_high\"],\n",
+    "            \"blur\": p[\"blur\"], \"blur_kernel\": p[\"blur_kernel\"], \"blur_sigma\": p[\"blur_sigma\"],\n",
+    "            \"interpolation\": p[\"interpolation\"],\n",
+    "            \"time_from\": t_from, \"time_to\": t_to,\n",
+    "        }"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Patches erzeugen und als Parquet schreiben\n",
+    "\n",
+    "Shards ≤ ~500 MB wie `data/train-00000-of-00012.parquet`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "SHARD_ROWS = 800\n",
+    "\n",
+    "tile_index = {t[\"tile_id\"]: i for i, t in enumerate(tiles)}\n",
+    "split_ds = {}\n",
+    "\n",
+    "for split in [\"train\", \"validation\", \"test\"]:\n",
+    "    split_tiles = [t for t in tiles if t[\"split\"] == split]\n",
+    "\n",
+    "    def gen(split_tiles=split_tiles):\n",
+    "        for t in tqdm(split_tiles, desc=f\"Patches {split}\", ncols=120):\n",
+    "            yield from tile_rows(t, tile_index[t[\"tile_id\"]])\n",
+    "\n",
+    "    ds = Dataset.from_generator(gen, features=features, cache_dir=LOCAL_CACHE)\n",
+    "    ds = ds.shuffle(seed=config[\"SEED\"]).flatten_indices()\n",
+    "    split_ds[split] = ds\n",
+    "    print(f\"{split}: {len(ds)} patches from {len(split_tiles)} tiles\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Alte Shards entfernen (sonst Mischung aus Läufen)\n",
+    "shutil.rmtree(os.path.join(LOCAL_OUT, \"data\"), ignore_errors=True)\n",
+    "\n",
+    "for split, ds in split_ds.items():\n",
+    "    n_shards = max(1, math.ceil(len(ds) / SHARD_ROWS))\n",
+    "    for i in range(n_shards):\n",
+    "        out = os.path.join(LOCAL_OUT, \"data\", f\"{split}-{i:05d}-of-{n_shards:05d}.parquet\")\n",
+    "        os.makedirs(os.path.dirname(out), exist_ok=True)\n",
+    "        ds.shard(n_shards, i, contiguous=True).to_parquet(out)\n",
+    "    print(f\"{split}: {n_shards} shards\")\n",
+    "\n",
+    "!du -sh {LOCAL_OUT}/data"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Auf Drive kopieren (ersetzt alten Stand)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "for f in glob(os.path.join(HF_DATA_DIR, \"*.parquet\")):\n",
+    "    os.remove(f)\n",
+    "for f in sorted(glob(os.path.join(LOCAL_OUT, \"data\", \"*.parquet\"))):\n",
+    "    shutil.copy(f, HF_DATA_DIR)\n",
+    "\n",
+    "tiles_df.to_csv(os.path.join(HF_META_DIR, \"tiles.csv\"), index=False)\n",
+    "print(sorted(os.listdir(HF_DATA_DIR)))"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 5. Checks"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "hf = load_dataset(\"parquet\", data_dir=HF_DATA_DIR, cache_dir=LOCAL_CACHE)\n",
+    "print(hf)"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Kachel-IDs disjunkt?\n",
+    "ids = {s: set(hf[s][\"tile_id\"]) for s in hf}\n",
+    "assert not (ids[\"train\"] & ids[\"validation\"]), \"train/val leak\"\n",
+    "assert not (ids[\"train\"] & ids[\"test\"]), \"train/test leak\"\n",
+    "assert not (ids[\"validation\"] & ids[\"test\"]), \"val/test leak\"\n",
+    "\n",
+    "# Patch-IDs eindeutig?\n",
+    "all_ids = sum((list(hf[s][\"patch_id\"]) for s in hf), [])\n",
+    "assert len(all_ids) == len(set(all_ids))\n",
+    "\n",
+    "# Geometrisch disjunkt (Patch-Ebene)?\n",
+    "boxes = {s: hf[s].select_columns([\"lon_min\", \"lat_min\", \"lon_max\", \"lat_max\"]).to_pandas().values for s in hf}\n",
+    "for a in hf:\n",
+    "    for b in hf:\n",
+    "        if a < b:\n",
+    "            A, Bx = boxes[a], boxes[b]\n",
+    "            ov = ((A[:, None, 0] < Bx[None, :, 2]) & (Bx[None, :, 0] < A[:, None, 2]) &\n",
+    "                  (A[:, None, 1] < Bx[None, :, 3]) & (Bx[None, :, 1] < A[:, None, 3])).sum()\n",
+    "            assert ov == 0, f\"{a}/{b}: {ov} overlapping patches\"\n",
+    "\n",
+    "x = hf[\"train\"][0]\n",
+    "print(x[\"lr\"].size, x[\"hr\"].size, np.array(x[\"ndvi\"]).shape, np.array(x[\"reflectance\"]).shape)\n",
+    "print(\"OK: splits disjoint by tile, patch id and footprint.\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 6. Statistiken und Abbildungen für die Dataset Card"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "SPLITS = [\"train\", \"validation\", \"test\"]\n",
+    "SPLIT_COLORS = {\"train\": \"#2a78d6\", \"validation\": \"#eb6834\", \"test\": \"#1baf7a\"}\n",
+    "SURFACE, INK, INK2, GRID = \"#fcfcfb\", \"#0b0b0b\", \"#52514e\", \"#e4e3df\"\n",
+    "\n",
+    "plt.rcParams.update({\n",
+    "    \"figure.facecolor\": SURFACE, \"axes.facecolor\": SURFACE, \"savefig.facecolor\": SURFACE,\n",
+    "    \"axes.edgecolor\": GRID, \"axes.labelcolor\": INK2, \"xtick.color\": INK2, \"ytick.color\": INK2,\n",
+    "    \"text.color\": INK, \"axes.grid\": True, \"grid.color\": GRID, \"grid.linewidth\": 0.6,\n",
+    "    \"axes.spines.top\": False, \"axes.spines.right\": False, \"font.size\": 11,\n",
+    "    \"axes.titleweight\": \"bold\", \"axes.titlesize\": 13, \"legend.frameon\": False,\n",
+    "})\n",
+    "\n",
+    "def save_fig(fig, name):\n",
+    "    path = os.path.join(HF_IMAGES_DIR, name)\n",
+    "    fig.savefig(path, dpi=160, bbox_inches=\"tight\")\n",
+    "    print(\"saved\", path)\n",
+    "\n",
+    "meta = {s: hf[s].select_columns([c for c in hf[s].column_names if c not in\n",
+    "        (\"lr\", \"hr\", \"ndvi\", \"scl\", \"reflectance\")]).to_pandas() for s in SPLITS}"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### LR/HR-Paare (Test-Split)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "plt.rcdefaults()\n",
+    "plt.rcParams.update({\"figure.facecolor\": \"white\", \"savefig.facecolor\": \"white\", \"font.size\": 12})\n",
+    "\n",
+    "def plain(axs):\n",
+    "    for a in np.ravel(axs):\n",
+    "        a.set_xticks([]); a.set_yticks([])\n",
+    "        for sp in a.spines.values():\n",
+    "            sp.set_visible(False)\n",
+    "\n",
+    "# Ausgewählte Test-Patches für die Dataset Card\n",
+    "PAIR_IDS = ['aoi_0424_p03', 'custom_0004_p02', 'aoi_0176_p10', 'aoi_0374_p02', 'custom_0002_p00', 'aoi_0261_p02']\n",
+    "DEG_ID = 'custom_0006_p07'\n",
+    "pid = list(meta[\"test\"].patch_id)\n",
+    "pair_idx = [pid.index(p) for p in PAIR_IDS]\n",
+    "\n",
+    "fig, axs = plt.subplots(2, len(pair_idx), figsize=(3.1 * len(pair_idx), 6.6))\n",
+    "for c, i in enumerate(pair_idx):\n",
+    "    r = hf[\"test\"][i]\n",
+    "    axs[0, c].imshow(r[\"lr\"], interpolation=\"nearest\")\n",
+    "    axs[1, c].imshow(r[\"hr\"])\n",
+    "    axs[0, c].set_title(r[\"aoi_name\"] or f\"{r['lat']:.2f}°N {r['lon']:.2f}°E\", fontsize=12)\n",
+    "axs[0, 0].set_ylabel(\"LR input (64×64)\", fontsize=13)\n",
+    "axs[1, 0].set_ylabel(\"HR target (256×256)\", fontsize=13)\n",
+    "plain(axs)\n",
+    "plt.tight_layout(pad=0.4)\n",
+    "save_fig(fig, \"psrgan_s2_pairs.png\"); plt.show()"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Land Cover (SCL) je Split (für README)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "SCL_NAMES = {2: \"Dark area\", 3: \"Cloud shadow\", 4: \"Vegetation\", 5: \"Not vegetated\", 6: \"Water\",\n",
+    "             7: \"Unclassified\", 8: \"Cloud (medium)\", 9: \"Cloud (high)\", 10: \"Thin cirrus\", 11: \"Snow/ice\"}\n",
+    "\n",
+    "scl_frac = {}\n",
+    "for s in SPLITS:\n",
+    "    counts = np.zeros(12)\n",
+    "    for batch in hf[s].select_columns([\"scl\"]).with_format(\"numpy\").iter(batch_size=256):\n",
+    "        counts += np.bincount(batch[\"scl\"].ravel(), minlength=12)[:12]\n",
+    "    scl_frac[s] = counts / counts.sum()\n",
+    "\n",
+    "for c in SCL_NAMES:\n",
+    "    print(f\"{SCL_NAMES[c]:15s}\", \"  \".join(f\"{s}: {scl_frac[s][c]:6.1%}\" for s in SPLITS))"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### NDVI je Split (für README)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "bins = np.linspace(-1, 1, 81)\n",
+    "ndvi_stats = {}\n",
+    "for s in SPLITS:\n",
+    "    h = np.zeros(len(bins) - 1)\n",
+    "    for batch in hf[s].select_columns([\"ndvi\"]).with_format(\"numpy\").iter(batch_size=256):\n",
+    "        h += np.histogram(batch[\"ndvi\"].ravel(), bins=bins)[0]\n",
+    "    h = h / h.sum() / np.diff(bins)\n",
+    "    ndvi_stats[s] = float(((bins[:-1] + bins[1:]) / 2 * h * np.diff(bins)).sum())\n",
+    "print(\"NDVI-Mittel:\", {s: round(v, 3) for s, v in ndvi_stats.items()})"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Degradations-Pipeline an einem Patch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Echter downsample_realistic-Aufruf, je Filter ein Beispiel (unten: Ausschnitt vergrößert)\n",
+    "src = pid.index(DEG_ID)\n",
+    "hr_img = np.asarray(hf[\"test\"][src][\"hr\"])[..., :3].astype(np.float32) / 255\n",
+    "targets = [(\"gaussian\", \"bicubic\"), (\"motion_h\", \"bilinear\"), (\"motion_v\", \"area\"), (\"defocus\", \"bicubic\")]\n",
+    "found, seed = {}, 0\n",
+    "while len(found) < len(targets) and seed < 5000:\n",
+    "    random.seed(seed)\n",
+    "    lr, p = downsample_realistic(hr_img.transpose(2, 0, 1), UPSCALE, return_params=True)\n",
+    "    if (p[\"blur\"], p[\"interpolation\"]) in targets:\n",
+    "        found.setdefault((p[\"blur\"], p[\"interpolation\"]), (lr, p))\n",
+    "    seed += 1\n",
+    "\n",
+    "BLUR_NAMES = {\"gaussian\": \"Gaussian\", \"motion_h\": \"Motion (horizontal)\", \"motion_v\": \"Motion (vertical)\", \"defocus\": \"Defocus\"}\n",
+    "keys = [t for t in targets if t in found]\n",
+    "# Ausschnitt mit den meisten Kanten (64×64 HR bzw. 16×16 LR)\n",
+    "g = hr_img.mean(-1)\n",
+    "gm = np.abs(np.diff(g, axis=0))[:, :-1] + np.abs(np.diff(g, axis=1))[:-1, :]\n",
+    "cy, cx = max(((y, x) for y in range(0, 193, 16) for x in range(0, 193, 16)),\n",
+    "             key=lambda q: gm[q[0]:q[0] + 64, q[1]:q[1] + 64].sum())\n",
+    "fig, axs = plt.subplots(2, 1 + len(keys), figsize=(3.3 * (1 + len(keys)), 7.2))\n",
+    "axs[0, 0].imshow(hr_img); axs[0, 0].set_title(\"HR (256×256)\", fontsize=12)\n",
+    "axs[0, 0].add_patch(plt.Rectangle((cx, cy), 64, 64, fill=False, edgecolor=\"white\", linewidth=1.5))\n",
+    "axs[1, 0].imshow(hr_img[cy:cy + 64, cx:cx + 64], interpolation=\"nearest\")\n",
+    "for c, key in enumerate(keys, start=1):\n",
+    "    lr, p = found[key]\n",
+    "    lr_img = np.clip(lr.transpose(1, 2, 0), 0, 1)\n",
+    "    axs[0, c].imshow(lr_img, interpolation=\"nearest\")\n",
+    "    axs[1, c].imshow(lr_img[cy // 4:(cy + 64) // 4, cx // 4:(cx + 64) // 4], interpolation=\"nearest\")\n",
+    "    blur = BLUR_NAMES[p[\"blur\"]] + f\" {p['blur_kernel']}×{p['blur_kernel']}\"\n",
+    "    if p[\"blur\"] == \"gaussian\":\n",
+    "        blur += f\", σ {p['blur_sigma']:.1f}\"\n",
+    "    axs[0, c].set_title(f\"{blur}\\n+ {p['interpolation']} ↓4\", fontsize=12)\n",
+    "axs[0, 0].set_ylabel(\"full patch\", fontsize=12)\n",
+    "axs[1, 0].set_ylabel(\"zoom (white box)\", fontsize=12)\n",
+    "plain(axs)\n",
+    "plt.tight_layout(pad=0.4)\n",
+    "save_fig(fig, \"psrgan_s2_degradation.png\"); plt.show()"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Kennzahlen (für README)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "def dir_bytes(pattern):\n",
+    "    return int(sum(os.path.getsize(f) for f in glob(pattern)))\n",
+    "\n",
+    "stats = {\n",
+    "    \"created\": time.strftime(\"%Y-%m-%d\"),\n",
+    "    \"region\": config[\"REGION\"],\n",
+    "    \"region_v1\": config[\"REGION_V1\"],\n",
+    "    \"time_window\": [config[\"TIME_FROM\"], config[\"TIME_TO\"]],\n",
+    "    \"patches_per_tile\": config[\"PATCHES_PER_TILE\"],\n",
+    "    \"n_tiles_total\": len(tiles),\n",
+    "    \"n_blocks\": len(block_groups),\n",
+    "    \"splits\": {},\n",
+    "}\n",
+    "for s in SPLITS:\n",
+    "    m = meta[s]\n",
+    "    stats[\"splits\"][s] = {\n",
+    "        \"num_examples\": len(hf[s]),\n",
+    "        \"num_tiles\": int(m.tile_id.nunique()),\n",
+    "        \"num_city_patches\": int((m.source == \"city\").sum()),\n",
+    "        \"num_bytes_parquet\": dir_bytes(os.path.join(HF_DATA_DIR, f\"{s}-*.parquet\")),\n",
+    "        \"num_bytes_arrow\": int(split_ds[s].data.nbytes),\n",
+    "        \"num_shards\": len(glob(os.path.join(HF_DATA_DIR, f\"{s}-*.parquet\"))),\n",
+    "        \"cloud_fraction_mean\": float(m.cloud_fraction.mean()),\n",
+    "        \"water_fraction_mean\": float(m.water_fraction.mean()),\n",
+    "        \"ndvi_mean\": ndvi_stats[s],\n",
+    "        \"scl_fraction\": {SCL_NAMES.get(c, str(c)): float(scl_frac[s][c]) for c in range(12) if scl_frac[s][c] > 0},\n",
+    "        \"lat_range\": [float(m.lat.min()), float(m.lat.max())],\n",
+    "        \"lon_range\": [float(m.lon.min()), float(m.lon.max())],\n",
+    "    }\n",
+    "stats[\"num_examples_total\"] = sum(v[\"num_examples\"] for v in stats[\"splits\"].values())\n",
+    "stats[\"download_size\"] = sum(v[\"num_bytes_parquet\"] for v in stats[\"splits\"].values())\n",
+    "\n",
+    "json.dump(stats, open(os.path.join(HF_META_DIR, \"stats.json\"), \"w\"), indent=2, ensure_ascii=False)\n",
+    "print(json.dumps(stats, indent=2, ensure_ascii=False))"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "!ls -la {HF_DATASET_DIR} {HF_DATA_DIR} {HF_IMAGES_DIR} {HF_META_DIR}"
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "colab": {
+   "machine_shape": "hm",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}

psrgan.py

@@ -0,0 +1,216 @@
+"""PSRGAN: probabilistic 4x super-resolution for Sentinel-2 RGB.
+
+Standalone model code (PyTorch only). The backbone is a copy of the RRDBNet
+architecture from BasicSR / Real-ESRGAN (BSD-3-Clause), with module names kept
+identical so the published weights load strictly.
+
+    from psrgan import PSRGAN
+    model = PSRGAN.from_pretrained("marklpd/PSRGAN").eval()
+    out = model.super_resolve(lr)   # lr: (B, 3, H, W) in [0, 1]
+    out["sr"], out["sigma_aleatoric"]
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import PyTorchModelHubMixin
+
+__all__ = ["PSRGAN", "RRDBNet", "BayesianConv2d"]
+
+
+# ---------------------------------------------------------------- backbone --
+class ResidualDenseBlock(nn.Module):
+    def __init__(self, num_feat: int = 64, num_grow_ch: int = 32):
+        super().__init__()
+        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
+        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
+        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+    def forward(self, x):
+        x1 = self.lrelu(self.conv1(x))
+        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
+        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
+        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
+        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
+        return x5 * 0.2 + x
+
+
+class RRDB(nn.Module):
+    def __init__(self, num_feat: int, num_grow_ch: int = 32):
+        super().__init__()
+        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
+        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
+
+    def forward(self, x):
+        out = self.rdb1(x)
+        out = self.rdb2(out)
+        out = self.rdb3(out)
+        return out * 0.2 + x
+
+
+class RRDBNet(nn.Module):
+    """RRDBNet (x4) with the same parameter names as basicsr.archs.rrdbnet_arch.RRDBNet."""
+
+    def __init__(self, num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32):
+        super().__init__()
+        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
+        self.body = nn.Sequential(*[RRDB(num_feat, num_grow_ch) for _ in range(num_block)])
+        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
+        # Unused by PSRGAN (the mu/s heads replace it); kept so Real-ESRGAN keys line up.
+        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
+        self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)
+
+
+# ------------------------------------------------------------------- heads --
+class BayesianConv2d(nn.Module):
+    """3x3 conv with a Gaussian over its weights (mean + log-variance)."""
+
+    def __init__(self, in_ch: int, out_ch: int, init_logvar: float = -10.0):
+        super().__init__()
+        conv = nn.Conv2d(in_ch, out_ch, 3, 1, 1)
+        self.weight_mu = nn.Parameter(conv.weight.data.clone())
+        self.weight_logvar = nn.Parameter(torch.full_like(conv.weight, init_logvar))
+        self.bias = conv.bias
+
+    def forward(self, x, sample: bool = True):
+        # Weights are only sampled in train mode, exactly as during training.
+        if self.training and sample:
+            eps = torch.randn_like(self.weight_mu)
+            weight = self.weight_mu + eps * torch.exp(0.5 * self.weight_logvar)
+        else:
+            weight = self.weight_mu
+        return F.conv2d(x, weight, self.bias, stride=1, padding=1)
+
+
+# ------------------------------------------------------------------- model --
+class PSRGAN(
+    nn.Module,
+    PyTorchModelHubMixin,
+    library_name="pytorch",
+    license="bsd-3-clause",
+    pipeline_tag="image-to-image",
+    tags=["super-resolution", "sentinel-2", "remote-sensing", "uncertainty", "esrgan"],
+):
+    """4x super-resolution with a per-pixel mean (mu) and log-variance (s) output.
+
+    forward(lr) returns (mu, s): mu in [0, 1] is the super-resolved image,
+    s = log(sigma^2) is the aleatoric uncertainty per pixel and channel.
+    """
+
+    def __init__(self, num_feat: int = 64, num_block: int = 23, num_grow_ch: int = 32,
+                 scale: int = 4, bayesian_heads: bool = True, logvar_min: float = -10.0,
+                 logvar_max: float = 10.0):
+        super().__init__()
+        assert scale == 4, "PSRGAN is trained for x4"
+        self.scale = scale
+        self.logvar_min, self.logvar_max = logvar_min, logvar_max
+        self.backbone = RRDBNet(3, 3, num_feat, num_block, num_grow_ch)
+        if bayesian_heads:
+            self.mu_head = BayesianConv2d(num_feat, 3)
+            self.s_head = BayesianConv2d(num_feat, 3)
+        else:
+            self.mu_head = nn.Conv2d(num_feat, 3, 3, 1, 1)
+            self.s_head = nn.Conv2d(num_feat, 3, 3, 1, 1)
+        self._swag = None
+
+    def features(self, lr):
+        b = self.backbone
+        feat = b.conv_first(lr)
+        feat = feat + b.conv_body(b.body(feat))
+        feat = F.leaky_relu(b.conv_up1(F.interpolate(feat, scale_factor=2, mode="nearest")), 0.2, inplace=True)
+        feat = F.leaky_relu(b.conv_up2(F.interpolate(feat, scale_factor=2, mode="nearest")), 0.2, inplace=True)
+        return F.leaky_relu(b.conv_hr(feat), 0.2, inplace=True)
+
+    def forward(self, lr, sample_bayes: bool = False):
+        feat = self.features(lr)
+        if isinstance(self.mu_head, BayesianConv2d):
+            mu = torch.sigmoid(self.mu_head(feat, sample=sample_bayes))
+            s = self.s_head(feat, sample=sample_bayes)
+        else:
+            mu = torch.sigmoid(self.mu_head(feat))
+            s = self.s_head(feat)
+        return mu, torch.clamp(s, self.logvar_min, self.logvar_max)
+
+    # --------------------------------------------------------------- SWAG --
+    def load_swag(self, path: str):
+        """Load SWAG moments (psrgan_x4_swag.safetensors) for epistemic uncertainty."""
+        from safetensors.torch import load_file
+
+        t = load_file(path)
+        self._swag = {
+            "mean": {k[5:]: v for k, v in t.items() if k.startswith("mean.")},
+            "var": {k[4:]: v for k, v in t.items() if k.startswith("var.")},
+        }
+        return self
+
+    @torch.no_grad()
+    def _set_swag_sample(self, generator=None):
+        for n, p in self.named_parameters():
+            if n in self._swag["mean"]:
+                m = self._swag["mean"][n].to(p.device)
+                v = self._swag["var"][n].to(p.device)
+                eps = torch.randn(m.shape, generator=generator, device="cpu").to(p.device)
+                p.copy_(m + v.sqrt() * eps)
+
+    # ---------------------------------------------------------- inference --
+    @torch.no_grad()
+    def super_resolve(self, lr, mc_samples: int = 0, swag_samples: int = 0, epistemic_weight: float = 10.0):
+        """Super-resolve a batch.
+
+        Args:
+            lr: (B, 3, H, W) float tensor in [0, 1], RGB = (B04, B03, B02) after a 1-99 % stretch.
+            mc_samples: Monte-Carlo samples of the Bayesian heads (epistemic, cheap).
+            swag_samples: SWAG weight samples (epistemic, needs load_swag()).
+            epistemic_weight: scale of the epistemic variance in sigma_total (10 as in the paper code).
+
+        Returns:
+            dict with sr (B,3,4H,4W), sigma_aleatoric, and if requested sigma_epistemic and sigma_total.
+        """
+        was_training = self.training
+        self.eval()
+        mu, s = self(lr, sample_bayes=False)
+        out = {"sr": mu.clamp(0, 1), "sigma_aleatoric": torch.exp(0.5 * s)}
+
+        mus = []
+        if mc_samples > 0:
+            self.mu_head.train()
+            self.s_head.train()
+            mus += [self(lr, sample_bayes=True)[0] for _ in range(mc_samples)]
+            self.mu_head.eval()
+            self.s_head.eval()
+        if swag_samples > 0:
+            if self._swag is None:
+                raise RuntimeError("call load_swag() first")
+            backup = {k: v.clone() for k, v in self.state_dict().items()}
+            for _ in range(swag_samples):
+                self._set_swag_sample()
+                mus.append(self(lr, sample_bayes=False)[0])
+            self.load_state_dict(backup)
+        if mus:
+            var_epi = torch.stack(mus).var(0, unbiased=False)
+            out["sigma_epistemic"] = var_epi.sqrt()
+            out["sigma_total"] = torch.sqrt(out["sigma_aleatoric"] ** 2 + epistemic_weight * var_epi + 1e-8)
+
+        self.train(was_training)
+        return out
+
+
+def stretch(rgb_reflectance, low: float = 1.0, high: float = 99.0):
+    """The 1-99 % percentile stretch used to build the training data.
+
+    rgb_reflectance: (..., 3, H, W) Sentinel-2 L2A reflectance (B04, B03, B02), any scale.
+    Percentiles are computed over all channels and pixels of each image, like the dataset.
+    """
+    x = torch.as_tensor(rgb_reflectance, dtype=torch.float32)
+    flat = x.flatten(-3)
+    p_low = torch.quantile(flat, low / 100, dim=-1, keepdim=True)[..., None, None]
+    p_high = torch.quantile(flat, high / 100, dim=-1, keepdim=True)[..., None, None]
+    return ((x - p_low) / (p_high - p_low + 1e-8)).clamp(0, 1)

psrgan_x4_psnr.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a68bf9b007ddb4c3c172e5cca60b26e3f6acd68931adb6d472822e2ad36535c
+size 66892468

psrgan_x4_swag.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9d51a3cbc734f820875320f20e434abb26c2e15becaf6365fc945999463fa0b6
+size 133792144