| """Compact, trainable SkySense reproduction for multi-modal remote sensing data.""" |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class SpatialEncoder(nn.Module): |
| def __init__(self, in_channels, embed_dim, patch_size): |
| super().__init__() |
| self.projection = nn.Sequential( |
| nn.Conv2d(in_channels, embed_dim, patch_size, patch_size), |
| nn.GELU(), |
| nn.Conv2d(embed_dim, embed_dim, 3, padding=1), |
| nn.GELU(), |
| ) |
|
|
| def forward(self, images): |
| batch, time, channels, height, width = images.shape |
| features = self.projection(images.reshape(batch * time, channels, height, width)) |
| _, dim, out_height, out_width = features.shape |
| return features.reshape(batch, time, dim, out_height, out_width) |
|
|
|
|
| class SkySense(nn.Module): |
| """Factorized spatial-temporal encoder with geo-context prototypes.""" |
|
|
| def __init__( |
| self, |
| hr_channels=3, |
| s2_channels=10, |
| s1_channels=2, |
| embed_dim=32, |
| hr_patch_size=16, |
| s2_patch_size=8, |
| s1_patch_size=8, |
| temporal_depth=2, |
| temporal_heads=4, |
| num_regions=16, |
| prototypes_per_region=4, |
| num_classes=6, |
| ): |
| super().__init__() |
| self.num_regions = num_regions |
| self.hr_encoder = SpatialEncoder(hr_channels, embed_dim, hr_patch_size) |
| self.s2_encoder = SpatialEncoder(s2_channels, embed_dim, s2_patch_size) |
| self.s1_encoder = SpatialEncoder(s1_channels, embed_dim, s1_patch_size) |
| self.date_embedding = nn.Embedding(366, embed_dim) |
| self.modality_embedding = nn.Parameter(torch.zeros(3, embed_dim)) |
| self.fusion_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| layer = nn.TransformerEncoderLayer( |
| d_model=embed_dim, |
| nhead=temporal_heads, |
| dim_feedforward=embed_dim * 4, |
| dropout=0.0, |
| activation="gelu", |
| batch_first=True, |
| norm_first=True, |
| ) |
| self.temporal_fusion = nn.TransformerEncoder(layer, temporal_depth) |
| modality_layer = nn.TransformerEncoderLayer( |
| d_model=embed_dim, nhead=temporal_heads, dim_feedforward=embed_dim * 4, |
| dropout=0.0, activation="gelu", batch_first=True, norm_first=True, |
| ) |
| self.modality_fusion = nn.TransformerEncoder(modality_layer, 1) |
| self.modality_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) |
| self.prototypes = nn.Parameter( |
| torch.randn(num_regions, prototypes_per_region, embed_dim) * 0.02 |
| ) |
| self.decoder = nn.Sequential( |
| nn.Conv2d(embed_dim * 2, embed_dim, 3, padding=1), |
| nn.GELU(), |
| nn.Conv2d(embed_dim, num_classes, 1), |
| ) |
| nn.init.normal_(self.date_embedding.weight, std=0.02) |
| nn.init.normal_(self.modality_embedding, std=0.02) |
| nn.init.normal_(self.fusion_token, std=0.02) |
| nn.init.normal_(self.modality_token, std=0.02) |
|
|
| def _add_context(self, features, dates, modality_index): |
| if dates.dtype != torch.long: |
| raise TypeError(f"dates must use torch.int64, got {dates.dtype}") |
| if dates.shape != features.shape[:2]: |
| raise ValueError(f"dates shape {tuple(dates.shape)} does not match image batch/time {tuple(features.shape[:2])}") |
| if torch.any((dates < 0) | (dates > 364)): |
| raise ValueError("dates must contain day-of-year values in [0, 364]") |
| date_context = self.date_embedding(dates).unsqueeze(-1).unsqueeze(-1) |
| modality = self.modality_embedding[modality_index].view(1, 1, -1, 1, 1) |
| return features + date_context + modality |
|
|
| def encode_modalities(self, hr, s2, s1, dates_hr, dates_s2, dates_s1): |
| return ( |
| self._add_context(self.hr_encoder(hr), dates_hr, 0), |
| self._add_context(self.s2_encoder(s2), dates_s2, 1), |
| self._add_context(self.s1_encoder(s1), dates_s1, 2), |
| ) |
|
|
| def _aggregate_time(self, features): |
| batch, time, dim, height, width = features.shape |
| sequence = features.permute(0, 3, 4, 1, 2).reshape(-1, time, dim) |
| token = self.fusion_token.expand(sequence.shape[0], -1, -1) |
| fused = self.temporal_fusion(torch.cat([token, sequence], dim=1))[:, 0] |
| return fused.reshape(batch, height, width, dim).permute(0, 3, 1, 2) |
|
|
| def forward(self, hr, s2, s1, dates_hr, dates_s2, dates_s1, region): |
| if region.dtype != torch.long: |
| raise TypeError(f"region must use torch.int64, got {region.dtype}") |
| if region.shape != (hr.shape[0],): |
| raise ValueError(f"region must have shape [{hr.shape[0]}], got {tuple(region.shape)}") |
| if torch.any((region < 0) | (region >= self.num_regions)): |
| raise ValueError(f"region IDs must be in [0, {self.num_regions - 1}]") |
| modality_features = self.encode_modalities(hr, s2, s1, dates_hr, dates_s2, dates_s1) |
| aggregated = [self._aggregate_time(feature) for feature in modality_features] |
| target_size = aggregated[0].shape[-2:] |
| aligned = [aggregated[0]] + [ |
| F.interpolate(feature, size=target_size, mode="bilinear", align_corners=False) |
| for feature in aggregated[1:] |
| ] |
| batch, dim, out_height, out_width = aligned[0].shape |
| modalities = torch.stack(aligned, dim=1).permute(0, 3, 4, 1, 2).reshape(-1, 3, dim) |
| token = self.modality_token.expand(modalities.shape[0], -1, -1) |
| fused = self.modality_fusion(torch.cat([token, modalities], dim=1))[:, 0] |
| fused = fused.reshape(batch, out_height, out_width, dim) |
|
|
| regional_prototypes = self.prototypes[region] |
| query = F.normalize(fused, dim=-1) |
| keys = F.normalize(regional_prototypes, dim=-1) |
| attention = torch.einsum("bhwd,bpd->bhwp", query, keys).softmax(dim=-1) |
| geo_context = torch.einsum("bhwp,bpd->bhwd", attention, regional_prototypes) |
| output = torch.cat([fused, geo_context], dim=-1).permute(0, 3, 1, 2) |
| logits = self.decoder(output) |
| logits = F.interpolate(logits, size=hr.shape[-2:], mode="bilinear", align_corners=False) |
| return {"logits": logits, "features": modality_features, "fused": fused} |
|
|
| @staticmethod |
| def cross_modal_alignment_loss(features): |
| pooled = [F.normalize(feature.mean(dim=(1, 3, 4)), dim=-1) for feature in features] |
| losses = [1.0 - (pooled[i] * pooled[j]).sum(dim=-1).mean() for i in range(3) for j in range(i + 1, 3)] |
| return torch.stack(losses).mean() |
|
|
