""" GeoMTConvNeXt — Geospatial Multi-Task ConvNeXt =============================================== Multi-task geospatial prediction model for the embed2heights benchmark (ESA/ITU GeoFM challenge, Belgium/Netherlands). Architecture ------------ Backbone : ConvNeXt-Tiny, ImageNet-1K pretrained and fine-tuned. A learned 192→3 adapter projects the stacked AE+Tessera pixel embeddings to pseudo-RGB so the pretrained encoder can be reused. Fusion : TerraMind S1/S2 and THOR S1/S2 patch-context embeddings (4 × 768 channels at 16×16) are compressed and injected at the ConvNeXt bottleneck, adding global geographic context. Decoder : 4-level U-Net with skip connections (768→384→192→96 channels). High-res : A parallel path from the raw 192ch pixel embeddings bypasses the encoder entirely and is merged back at full 256×256 resolution, recovering fine-grained building edge detail. Heads : (1) 3-class cover segmentation (building / vegetation / water) with sigmoid activation. (2) Height regression via ordinal soft-expectation over 64 bins spanning [0, HMAX] metres — softmax followed by a dot product with bin centres. Parameters : 52 M Val score : 0.468 (fold 0, competition metric) Input dict keys --------------- alphaearth_emb : [B, 64, 256, 256] AlphaEarth pixel embeddings tessera_emb : [B, 128, 256, 256] Tessera pixel embeddings terramind_s1_emb : [B, 768, 16, 16] TerraMind Sentinel-1 patch context terramind_s2_emb : [B, 768, 16, 16] TerraMind Sentinel-2 patch context thor_s1_emb : [B, 768, 16, 16] THOR Sentinel-1 patch context thor_s2_emb : [B, 768, 16, 16] THOR Sentinel-2 patch context Output ------ out : [B, 4, 256, 256] — channels 0-2: cover ∈ [0,1]; channel 3: height (m) h_logits : [B, N_BINS, 256, 256] — raw height bin logits seg_logits : [B, 3, 256, 256] — raw cover logits aux : [B, 3, 64, 64] — deep supervision cover (training only) """ import torch import torch.nn as nn import torch.nn.functional as F N_BINS = 64 HMAX = 50.0 # max height in metres HMAX_NORM = HMAX / 30.0 # normalised ceiling used during training # ── Building blocks ─────────────────────────────────────────────────────────── def _gn(c: int) -> nn.GroupNorm: return nn.GroupNorm(min(32, c), c) class ResBlock(nn.Module): def __init__(self, cin: int, cout: int, dilation: int = 1): super().__init__() self.conv = nn.Sequential( nn.Conv2d(cin, cout, 3, padding=dilation, dilation=dilation, bias=False), _gn(cout), nn.GELU(), nn.Conv2d(cout, cout, 3, padding=1, bias=False), _gn(cout), ) self.skip = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity() self.act = nn.GELU() def forward(self, x: torch.Tensor) -> torch.Tensor: return self.act(self.conv(x) + self.skip(x)) def _up(cin: int, cout: int) -> nn.ConvTranspose2d: return nn.ConvTranspose2d(cin, cout, 2, stride=2) # ── GeoMTConvNeXt ───────────────────────────────────────────────────────────── class GeoMTConvNeXt(nn.Module): """ Geospatial Multi-Task ConvNeXt. Takes a dict of multi-source satellite embeddings and returns per-pixel predictions for building cover, vegetation cover, water cover, and height. """ def __init__(self, base: int = 64, n_bins: int = N_BINS, pretrained: bool = True): super().__init__() from torchvision.models import convnext_tiny, ConvNeXt_Tiny_Weights weights = ConvNeXt_Tiny_Weights.IMAGENET1K_V1 if pretrained else None feats = convnext_tiny(weights=weights).features # ConvNeXt-Tiny encoder stages: 96 @/4 · 192 @/8 · 384 @/16 · 768 @/32 self.enc1 = feats[0:2] self.enc2 = feats[2:4] self.enc3 = feats[4:6] self.enc4 = feats[6:8] C1, C2, C3, C4 = 96, 192, 384, 768 # Adapter: stack AE (64ch) + Tessera (128ch) → pseudo-RGB for ConvNeXt self.adapter = nn.Sequential( nn.Conv2d(192, 64, 3, padding=1), _gn(64), nn.GELU(), nn.Conv2d(64, 3, 1), ) self.register_buffer( "imnet_mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)) self.register_buffer( "imnet_std", torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) # Patch-context encoder: TerraMind S1/S2 + THOR S1/S2 → injected at bottleneck self.patch_enc = nn.Sequential( nn.Conv2d(3072, 256, 1), _gn(256), nn.GELU(), ResBlock(256, 256)) self.bott = ResBlock(C4 + 256, C4) # U-Net decoder self.u3 = _up(C4, C3); self.d3 = ResBlock(C3 + C3, C3) self.u2 = _up(C3, C2); self.d2 = ResBlock(C2 + C2, C2) self.u1 = _up(C2, C1); self.d1 = ResBlock(C1 + C1, base) self.aux_seg = nn.Conv2d(base, 3, 1) # deep supervision @64px # High-res path: bypasses encoder, preserves 1-3px building detail self.hires = nn.Sequential( nn.Conv2d(192, 64, 3, padding=1), _gn(64), nn.GELU(), ResBlock(64, 64)) self.up_to_full= nn.Sequential( _up(base, base), _gn(base), nn.GELU(), _up(base, 64), _gn(64), nn.GELU()) self.fuse = ResBlock(64 + 64, base) # Cover head (building / vegetation / water) self.seg_dec = nn.Sequential( ResBlock(base, base), nn.Conv2d(base, base, 3, padding=1), _gn(base), nn.GELU(), ) self.seg_head = nn.Conv2d(base, 3, 1) # Height head: ordinal soft-expectation over n_bins bins self.hgt_dec = nn.Sequential( ResBlock(base + base, base), nn.Conv2d(base, base, 3, padding=1), _gn(base), nn.GELU(), ) self.hgt_head = nn.Conv2d(base, n_bins, 1) self.n_bins = n_bins self.register_buffer( "bin_centers", torch.linspace(0.0, HMAX_NORM, n_bins)) def forward(self, batch: dict) -> tuple: # Pixel embeddings: AE (64ch) + Tessera (128ch) → [B, 192, 256, 256] pix = torch.cat([batch["alphaearth_emb"], batch["tessera_emb"]], dim=1) # Patch context: 4 × TerraMind/THOR streams → [B, 3072, 16, 16] patch = torch.cat([batch["terramind_s1_emb"], batch["terramind_s2_emb"], batch["thor_s1_emb"], batch["thor_s2_emb"]], dim=1) # Adapt pixel embeddings to pseudo-RGB and normalise to ImageNet range x = self.adapter(pix) x = (torch.sigmoid(x) - self.imnet_mean) / self.imnet_std # ConvNeXt-Tiny encoder c1 = self.enc1(x) c2 = self.enc2(c1) c3 = self.enc3(c2) c4 = self.enc4(c3) # Inject patch context at bottleneck pe = self.patch_enc(patch) pe = F.interpolate(pe, size=c4.shape[-2:], mode="bilinear", align_corners=False) b = self.bott(torch.cat([c4, pe], dim=1)) # U-Net decoder x = self.d3(torch.cat([self.u3(b), c3], dim=1)) x = self.d2(torch.cat([self.u2(x), c2], dim=1)) x = self.d1(torch.cat([self.u1(x), c1], dim=1)) # [B, base, 64, 64] aux = self.aux_seg(x) # High-res path fused at full resolution hi = self.hires(pix) # [B, 64, 256, 256] xf = self.up_to_full(x) # [B, 64, 256, 256] feat = self.fuse(torch.cat([xf, hi], dim=1)) # [B, base, 256, 256] # Cover head sfeat = self.seg_dec(feat) seg_logits = self.seg_head(sfeat) cover = torch.sigmoid(seg_logits) # Height head — ordinal soft-expectation hfeat = self.hgt_dec(torch.cat([feat, sfeat], dim=1)) h_logits = self.hgt_head(hfeat) p = h_logits.float().softmax(dim=1) height = (p * self.bin_centers.view(1, -1, 1, 1)).sum(dim=1, keepdim=True) height = height * 30.0 # denormalise to metres out = torch.cat([cover, height], dim=1) # [B, 4, 256, 256] return out, h_logits, seg_logits, aux