OneScience-Group/Scale-MAE

1"""Scale-aware masked autoencoder for paired remote-sensing resolutions."""2import math3import torch4from torch import nn5from torch.nn import functional as F6 7 8def sincos_2d(height, width, dim, device, dtype, scale=1.0):9    if dim % 4:10        raise ValueError("embedding dimension must be divisible by four")11    y, x = torch.meshgrid(torch.arange(height, device=device, dtype=dtype),12                          torch.arange(width, device=device, dtype=dtype), indexing="ij")13    quarter = dim // 414    omega = torch.exp(-math.log(10000.0) * torch.arange(quarter, device=device, dtype=dtype) /15                      max(quarter - 1, 1))16    x, y = x.flatten()[:, None] * scale, y.flatten()[:, None] * scale17    return torch.cat((torch.sin(x * omega), torch.cos(x * omega),18                      torch.sin(y * omega), torch.cos(y * omega)), dim=1)19 20 21class ScaleMAE(nn.Module):22    def __init__(self, input_size=32, target_size=64, patch_size=4, in_channels=3,23                 embed_dim=64, encoder_depth=2, encoder_heads=4, decoder_dim=48,24                 decoder_depth=1, decoder_heads=4, mask_ratio=0.75, reference_gsd=1.0,25                 blur_kernel=5, band_config=None, **legacy):26        super().__init__()27        input_size = legacy.get("image_size", input_size)28        self.input_size, self.target_size, self.patch_size = input_size, target_size, patch_size29        self.in_channels, self.embed_dim, self.mask_ratio = in_channels, embed_dim, mask_ratio30        self.reference_gsd, self.blur_kernel = reference_gsd, blur_kernel31        if input_size % patch_size or target_size % patch_size:32            raise ValueError("input_size and target_size must be divisible by patch_size")33        self.input_grid = input_size // patch_size34        self.target_grid = target_size // patch_size35        self.num_patches = self.input_grid ** 236        patch_dim = in_channels * patch_size ** 237        self.patch_embed = nn.Linear(patch_dim, embed_dim)38        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))39        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))40        enc = nn.TransformerEncoderLayer(embed_dim, encoder_heads, embed_dim * 4,41                                         dropout=0, activation="gelu", batch_first=True, norm_first=True)42        self.encoder, self.encoder_norm = nn.TransformerEncoder(enc, encoder_depth), nn.LayerNorm(embed_dim)43        self.decoder_input = nn.Linear(embed_dim, decoder_dim)44        dec = nn.TransformerEncoderLayer(decoder_dim, decoder_heads, decoder_dim * 4,45                                         dropout=0, activation="gelu", batch_first=True, norm_first=True)46        self.decoder, self.decoder_norm = nn.TransformerEncoder(dec, decoder_depth), nn.LayerNorm(decoder_dim)47        self.fpn = nn.Sequential(nn.Conv2d(decoder_dim, decoder_dim, 3, padding=1), nn.GELU(),48                                 nn.Conv2d(decoder_dim, decoder_dim, 3, padding=1), nn.GELU())49        self.low_head = nn.Conv2d(decoder_dim, in_channels, 1)50        self.high_head = nn.Conv2d(decoder_dim, in_channels, 1)51        self.band_config = band_config or {"low": {"kernel": blur_kernel}, "high": {"residual": True}}52        nn.init.normal_(self.cls_token, std=.02); nn.init.normal_(self.mask_token, std=.02)53 54    def patchify(self, images):55        b, c, h, w = images.shape; p = self.patch_size56        if (c, h, w) != (self.in_channels, self.input_size, self.input_size):57            raise ValueError(f"expected input {(self.in_channels, self.input_size, self.input_size)}, got {(c,h,w)}")58        return images.reshape(b, c, h//p, p, w//p, p).permute(0,2,4,1,3,5).reshape(b, -1, c*p*p)59 60    def bandpass_targets(self, target):61        k = int(self.band_config.get("low", {}).get("kernel", self.blur_kernel))62        if k < 3 or k % 2 == 0: raise ValueError("low-frequency kernel must be odd and at least three")63        low = F.avg_pool2d(target, k, 1, k//2, count_include_pad=False)64        return low, target - low65 66    def _positions(self, gsd, grid, dim):67        scales = gsd.to(dtype=self.cls_token.dtype).flatten() / self.reference_gsd68        base = sincos_2d(grid, grid, dim, gsd.device, self.cls_token.dtype)69        return base.unsqueeze(0) * scales[:, None, None]70 71    def forward(self, images, gsd, target=None, mask_ratio=None, return_features=True):72        if target is None: target = F.interpolate(images, (self.target_size, self.target_size), mode="bilinear", align_corners=False)73        if target.shape[-2:] != (self.target_size, self.target_size):74            raise ValueError("target resolution does not match target_size")75        if target.shape[:2] != (images.shape[0], self.in_channels):76            raise ValueError("target batch or channel dimensions do not match images")77        b = images.shape[0]; patches = self.patchify(images); n = patches.shape[1]78        ratio = self.mask_ratio if mask_ratio is None else float(mask_ratio)79        if not 0 <= ratio < 1: raise ValueError("mask_ratio must be in [0, 1)")80        noise = torch.rand(b, n, device=images.device)81        ids_shuffle = noise.argsort(dim=1); ids_restore = ids_shuffle.argsort(dim=1)82        keep = max(1, int(n * (1 - ratio))); ids_keep = ids_shuffle[:, :keep]83        visible = torch.gather(self.patch_embed(patches), 1, ids_keep[..., None].expand(-1, -1, self.embed_dim))84        visible = visible + torch.gather(self._positions(gsd, self.input_grid, self.embed_dim), 1,85                                         ids_keep[..., None].expand(-1, -1, self.embed_dim))86        encoded = self.encoder_norm(self.encoder(torch.cat((self.cls_token.expand(b,-1,-1), visible), 1)))87        cls, visible_encoded = encoded[:, :1], encoded[:, 1:]88        decoded = self.decoder_input(visible_encoded)89        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)90        full = full + self._positions(gsd, self.input_grid, full.shape[-1])91        full = self.decoder_norm(self.decoder(full))92        fmap = full.transpose(1, 2).reshape(b, -1, self.input_grid, self.input_grid)93        fmap = F.interpolate(fmap, (self.target_grid, self.target_grid), mode="bilinear", align_corners=False)94        fmap = self.fpn(fmap)95        low = F.interpolate(self.low_head(fmap), (self.target_size, self.target_size), mode="bilinear", align_corners=False)96        high = F.interpolate(self.high_head(fmap), (self.target_size, self.target_size), mode="bilinear", align_corners=False)97        low_target, high_target = self.bandpass_targets(target)98        mask = torch.ones(b, n, dtype=torch.bool, device=images.device); mask.scatter_(1, ids_keep, False)99        pixel_mask = F.interpolate(mask.view(b, 1, self.input_grid, self.input_grid).float(),100                                   (self.target_size, self.target_size), mode="nearest")101        masked = pixel_mask.sum().clamp_min(1.0) * self.in_channels102        low_loss = ((low-low_target).square() * pixel_mask).sum() / masked103        high_loss = ((high-high_target).square() * pixel_mask).sum() / masked104        return {"loss": low_loss + high_loss, "low_loss": low_loss, "high_loss": high_loss,105                "reconstruction": low + high, "low_reconstruction": low, "high_reconstruction": high,106                "low_target": low_target, "high_target": high_target, "mask": mask,107                "ids_restore": ids_restore, "features": cls[:, 0] if return_features else None}108