| """Scale-aware masked autoencoder for paired remote-sensing resolutions.""" |
| import math |
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| def sincos_2d(height, width, dim, device, dtype, scale=1.0): |
| if dim % 4: |
| raise ValueError("embedding dimension must be divisible by four") |
| y, x = torch.meshgrid(torch.arange(height, device=device, dtype=dtype), |
| torch.arange(width, device=device, dtype=dtype), indexing="ij") |
| quarter = dim // 4 |
| omega = torch.exp(-math.log(10000.0) * torch.arange(quarter, device=device, dtype=dtype) / |
| max(quarter - 1, 1)) |
| x, y = x.flatten()[:, None] * scale, y.flatten()[:, None] * scale |
| return torch.cat((torch.sin(x * omega), torch.cos(x * omega), |
| torch.sin(y * omega), torch.cos(y * omega)), dim=1) |
|
|
|
|
| class ScaleMAE(nn.Module): |
| def __init__(self, input_size=32, target_size=64, patch_size=4, in_channels=3, |
| embed_dim=64, encoder_depth=2, encoder_heads=4, decoder_dim=48, |
| decoder_depth=1, decoder_heads=4, mask_ratio=0.75, reference_gsd=1.0, |
| blur_kernel=5, band_config=None, **legacy): |
| super().__init__() |
| input_size = legacy.get("image_size", input_size) |
| self.input_size, self.target_size, self.patch_size = input_size, target_size, patch_size |
| self.in_channels, self.embed_dim, self.mask_ratio = in_channels, embed_dim, mask_ratio |
| self.reference_gsd, self.blur_kernel = reference_gsd, blur_kernel |
| if input_size % patch_size or target_size % patch_size: |
| raise ValueError("input_size and target_size must be divisible by patch_size") |
| self.input_grid = input_size // patch_size |
| self.target_grid = target_size // patch_size |
| self.num_patches = self.input_grid ** 2 |
| patch_dim = in_channels * patch_size ** 2 |
| self.patch_embed = nn.Linear(patch_dim, embed_dim) |
| self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim)) |
| enc = nn.TransformerEncoderLayer(embed_dim, encoder_heads, embed_dim * 4, |
| dropout=0, activation="gelu", batch_first=True, norm_first=True) |
| self.encoder, self.encoder_norm = nn.TransformerEncoder(enc, encoder_depth), nn.LayerNorm(embed_dim) |
| self.decoder_input = nn.Linear(embed_dim, decoder_dim) |
| dec = nn.TransformerEncoderLayer(decoder_dim, decoder_heads, decoder_dim * 4, |
| dropout=0, activation="gelu", batch_first=True, norm_first=True) |
| self.decoder, self.decoder_norm = nn.TransformerEncoder(dec, decoder_depth), nn.LayerNorm(decoder_dim) |
| self.fpn = nn.Sequential(nn.Conv2d(decoder_dim, decoder_dim, 3, padding=1), nn.GELU(), |
| nn.Conv2d(decoder_dim, decoder_dim, 3, padding=1), nn.GELU()) |
| self.low_head = nn.Conv2d(decoder_dim, in_channels, 1) |
| self.high_head = nn.Conv2d(decoder_dim, in_channels, 1) |
| self.band_config = band_config or {"low": {"kernel": blur_kernel}, "high": {"residual": True}} |
| nn.init.normal_(self.cls_token, std=.02); nn.init.normal_(self.mask_token, std=.02) |
|
|
| def patchify(self, images): |
| b, c, h, w = images.shape; p = self.patch_size |
| if (c, h, w) != (self.in_channels, self.input_size, self.input_size): |
| raise ValueError(f"expected input {(self.in_channels, self.input_size, self.input_size)}, got {(c,h,w)}") |
| return images.reshape(b, c, h//p, p, w//p, p).permute(0,2,4,1,3,5).reshape(b, -1, c*p*p) |
|
|
| def bandpass_targets(self, target): |
| k = int(self.band_config.get("low", {}).get("kernel", self.blur_kernel)) |
| if k < 3 or k % 2 == 0: raise ValueError("low-frequency kernel must be odd and at least three") |
| low = F.avg_pool2d(target, k, 1, k//2, count_include_pad=False) |
| return low, target - low |
|
|
| def _positions(self, gsd, grid, dim): |
| scales = gsd.to(dtype=self.cls_token.dtype).flatten() / self.reference_gsd |
| base = sincos_2d(grid, grid, dim, gsd.device, self.cls_token.dtype) |
| return base.unsqueeze(0) * scales[:, None, None] |
|
|
| def forward(self, images, gsd, target=None, mask_ratio=None, return_features=True): |
| if target is None: target = F.interpolate(images, (self.target_size, self.target_size), mode="bilinear", align_corners=False) |
| if target.shape[-2:] != (self.target_size, self.target_size): |
| raise ValueError("target resolution does not match target_size") |
| if target.shape[:2] != (images.shape[0], self.in_channels): |
| raise ValueError("target batch or channel dimensions do not match images") |
| b = images.shape[0]; patches = self.patchify(images); n = patches.shape[1] |
| ratio = self.mask_ratio if mask_ratio is None else float(mask_ratio) |
| if not 0 <= ratio < 1: raise ValueError("mask_ratio must be in [0, 1)") |
| noise = torch.rand(b, n, device=images.device) |
| ids_shuffle = noise.argsort(dim=1); ids_restore = ids_shuffle.argsort(dim=1) |
| keep = max(1, int(n * (1 - ratio))); ids_keep = ids_shuffle[:, :keep] |
| visible = torch.gather(self.patch_embed(patches), 1, ids_keep[..., None].expand(-1, -1, self.embed_dim)) |
| visible = visible + torch.gather(self._positions(gsd, self.input_grid, self.embed_dim), 1, |
| ids_keep[..., None].expand(-1, -1, self.embed_dim)) |
| encoded = self.encoder_norm(self.encoder(torch.cat((self.cls_token.expand(b,-1,-1), visible), 1))) |
| cls, visible_encoded = encoded[:, :1], encoded[:, 1:] |
| decoded = self.decoder_input(visible_encoded) |
| full = self.mask_token.to(decoded.dtype).expand(b, n, -1).clone(); full.scatter_(1, ids_keep[..., None].expand(-1,-1,decoded.shape[-1]), decoded) |
| full = full + self._positions(gsd, self.input_grid, full.shape[-1]) |
| full = self.decoder_norm(self.decoder(full)) |
| fmap = full.transpose(1, 2).reshape(b, -1, self.input_grid, self.input_grid) |
| fmap = F.interpolate(fmap, (self.target_grid, self.target_grid), mode="bilinear", align_corners=False) |
| fmap = self.fpn(fmap) |
| low = F.interpolate(self.low_head(fmap), (self.target_size, self.target_size), mode="bilinear", align_corners=False) |
| high = F.interpolate(self.high_head(fmap), (self.target_size, self.target_size), mode="bilinear", align_corners=False) |
| low_target, high_target = self.bandpass_targets(target) |
| mask = torch.ones(b, n, dtype=torch.bool, device=images.device); mask.scatter_(1, ids_keep, False) |
| pixel_mask = F.interpolate(mask.view(b, 1, self.input_grid, self.input_grid).float(), |
| (self.target_size, self.target_size), mode="nearest") |
| masked = pixel_mask.sum().clamp_min(1.0) * self.in_channels |
| low_loss = ((low-low_target).square() * pixel_mask).sum() / masked |
| high_loss = ((high-high_target).square() * pixel_mask).sum() / masked |
| return {"loss": low_loss + high_loss, "low_loss": low_loss, "high_loss": high_loss, |
| "reconstruction": low + high, "low_reconstruction": low, "high_reconstruction": high, |
| "low_target": low_target, "high_target": high_target, "mask": mask, |
| "ids_restore": ids_restore, "features": cls[:, 0] if return_features else None} |
|
|
