"""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)