marklpd/PSRGAN

Test patches

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.

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.58 (Spearman ρ = 0.65), and the 10% most uncertain pixels have 12.8× the error of the 10% most confident ones.
3. On 1,328 test patches from unseen regions, PSRGAN reaches an LPIPS of 0.197, against 0.353 for Real-ESRGAN x4plus and 0.604 for bicubic upsampling. It also beats Real-ESRGAN in PSNR (21.23 vs. 20.20 dB) and SSIM (0.556 vs. 0.469).
4. Two sets of weights. The final adversarial weights have the lowest LPIPS, the weights after the likelihood stage the highest PSNR (21.82 dB) and SSIM (0.593) of all four methods tested.
5. Trained on randomized degradations (Gaussian, motion or defocus blur, then bicubic, bilinear or area downsampling), so it learns actual 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).

Prediction, error, and uncertainty.

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, that only needs torch, huggingface_hub and safetensors.

pip install torch huggingface_hub safetensors

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
degradation.py	randomized degradation that made the LR images of the dataset

Evaluation

All numbers are on the test split of PSRGAN Sentinel-2 version 1.1 (1,328 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.50 ± 0.14	0.5364 ± 0.0064	0.6042 ± 0.0040	8.274	0.00323
Real-ESRGAN	20.20 ± 0.13	0.4692 ± 0.0057	0.3529 ± 0.0035	8.760	0.00399
PSRGAN (PSNR)	21.82 ± 0.15	0.5926 ± 0.0054	0.2160 ± 0.0023	4.909	0.00287
PSRGAN	21.23 ± 0.14	0.5556 ± 0.0053	0.1971 ± 0.0022	5.645	0.00222

PSRGAN beats Real-ESRGAN in PSNR on 98% of test patches and in LPIPS on 99.9%. Compared with the likelihood stage, the adversarial stage gives up 0.60 dB PSNR and changes LPIPS from 0.216 to 0.197 (lower is better). Bicubic upsampling keeps a higher PSNR than the final model on 74% of patches, because PSNR rewards smooth averages over sharp but slightly misplaced detail. LPIPS (lower is better), 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.23 ± 0.28	0.5953 ± 0.0201	0.2138 ± 0.0089	5.894	0.00325
PSRGAN	20.74 ± 0.27	0.5648 ± 0.0191	0.2000 ± 0.0096	6.763	0.00255

Test patch from Herzogenaurach.

Calibration

	ECE (σ vs. RMSE)	error within 1σ	error within 2σ
Aleatoric + SWAG	0.0098	79.3%	96.8%
Aleatoric + Bayesian heads	0.0164	82.5%	97.5%

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.57) than on edges (r = 0.40).

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.2% Alpine tiles (Switzerland, Austria, northern Italy, south-eastern France, southern Germany, Slovenia) from late September and early October 2025. Expect worse results on other climates, other seasons, deserts and especially dense city centres.
• The uncertainty maps rank pixels by expected error, but the absolute σ is too large: 79% 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.27 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}
}