model.py

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