| """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"] |
|
|
|
|
| |
| 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) |
| |
| self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1) |
| self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True) |
|
|
|
|
| |
| 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): |
| |
| 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) |
|
|
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| @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) |
|
|
