modeling_geotag.py

"""
Model definition + loading for the JKTSV DINOv3 geolocation regressor ("modelD").

Architecture (must match the training checkpoint exactly):

    DinoGeoRegressor
      ├── encoder : DinoLastBlockEncoder
      │              └── backbone : DINOv3 ViT-L/16 (frozen)
      │                 forward = run backbone, grab the LAST transformer block
      │                 output (B, 201, 1024) via a forward hook, then
      │                 flatten -> (B, 205824)
      └── head    : UNetMLPHead(embed_dim=205824, hidden_dim=512, out_dim=2)

The head regresses a *local flat-earth (x, y) offset in metres* relative to a
fixed Jakarta origin. `local_xy_to_lonlat` inverts that projection to recover
(lon, lat) degrees. These constants are baked into the trained weights — do not
change them for inference.

The published checkpoint bundles the full (frozen) backbone weights together
with the trained head, so loading needs only the DINOv3 *architecture* from
``torch.hub`` (``pretrained=False``) — no separate LVD-1689M download.
"""

from __future__ import annotations

import os
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F

# --- constants that are part of the trained model -----------------------------

DINOV3_REPO = "facebookresearch/dinov3"
BACKBONE_NAME = "dinov3_vitl16"
EMBED_DIM = 205824        # 201 tokens (1 CLS + 4 storage + 196 patch) * 1024
HIDDEN_DIM = 512
OUT_DIM = 2

# Local flat-earth projection origin (Jakarta city centre) used during training.
ORIGIN_LON = 106.828320
ORIGIN_LAT = -6.227468
EARTH_RADIUS_M = 6371000.0


# --- coordinate conversion -----------------------------------------------------

def local_xy_to_lonlat(xy_meters: torch.Tensor) -> torch.Tensor:
    """Invert the flat-earth projection used as the regression target.

    Args:
        xy_meters: (B, 2) tensor of [x (east), y (north)] in metres.

    Returns:
        (B, 2) tensor of [lon, lat] in degrees.
    """
    lat0_rad = torch.deg2rad(torch.tensor(ORIGIN_LAT, device=xy_meters.device))
    x = xy_meters[:, 0]
    y = xy_meters[:, 1]

    dlon_rad = x / (EARTH_RADIUS_M * torch.cos(lat0_rad))
    dlat_rad = y / EARTH_RADIUS_M

    lon = torch.rad2deg(dlon_rad) + ORIGIN_LON
    lat = torch.rad2deg(dlat_rad) + ORIGIN_LAT
    return torch.stack([lon, lat], dim=-1)


# --- modules -------------------------------------------------------------------

class UNetMLPHead(nn.Module):
    """U-shaped MLP with 1-D skip connections. Input (B, embed_dim) -> (B, out_dim)."""

    def __init__(self, embed_dim: int, hidden_dim: int, out_dim: int):
        super().__init__()
        self.enc1 = nn.Linear(embed_dim, hidden_dim)
        self.enc2 = nn.Linear(hidden_dim, hidden_dim)
        self.bottleneck = nn.Linear(hidden_dim, hidden_dim)
        self.dec2 = nn.Linear(hidden_dim * 2, hidden_dim)
        self.dec1 = nn.Linear(hidden_dim * 2, hidden_dim)
        self.out = nn.Linear(hidden_dim, out_dim)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        e1 = F.gelu(self.enc1(x))
        e2 = F.gelu(self.enc2(e1))
        b = F.gelu(self.bottleneck(e2))
        d2 = F.gelu(self.dec2(torch.cat([b, e2], dim=-1)))
        d1 = F.gelu(self.dec1(torch.cat([d2, e1], dim=-1)))
        return self.out(d1)


class DinoLastBlockEncoder(nn.Module):
    """Run a DINOv3 ViT and return the flattened token sequence of its last block.

    A forward hook captures the last transformer block output (B, N, C); the
    tokens are flattened to (B, N*C). The backbone is frozen.
    """

    def __init__(self, backbone: nn.Module):
        super().__init__()
        self.backbone = backbone
        self._last_block_out = None
        self.backbone.blocks[-1].register_forward_hook(self._hook)
        for p in self.backbone.parameters():
            p.requires_grad = False

    def _hook(self, module, inputs, output):
        self._last_block_out = output

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        self._last_block_out = None
        _ = self.backbone(x)
        feats = self._last_block_out[0]   # (B, N, C)
        return feats.flatten(start_dim=1)  # (B, N*C)


class DinoGeoRegressor(nn.Module):
    """Frozen DINOv3 encoder + trainable UNet-MLP regression head.

    forward(pixel_values) -> (B, 2) local (x, y) metres.
    predict_lonlat(pixel_values) -> (B, 2) [lon, lat] degrees.

    `pixel_values` must already be resized to 224x224 and ImageNet-normalised
    (see ``GeoTagPredictor`` / the transform in ``inference.py``).
    """

    def __init__(
        self,
        backbone: nn.Module,
        embed_dim: int = EMBED_DIM,
        hidden_dim: int = HIDDEN_DIM,
        out_dim: int = OUT_DIM,
    ):
        super().__init__()
        self.encoder = DinoLastBlockEncoder(backbone)
        self.head = UNetMLPHead(embed_dim, hidden_dim, out_dim)

    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
        feats = self.encoder(pixel_values)
        return self.head(feats)

    @torch.no_grad()
    def predict_lonlat(self, pixel_values: torch.Tensor) -> torch.Tensor:
        return local_xy_to_lonlat(self.forward(pixel_values))

    # -- construction helpers --------------------------------------------------

    @staticmethod
    def build_backbone(device: str | torch.device = "cpu") -> nn.Module:
        """Instantiate the DINOv3 ViT-L/16 architecture (no pretrained download)."""
        backbone = torch.hub.load(
            DINOV3_REPO, BACKBONE_NAME, pretrained=False, trust_repo=True
        )
        return backbone.to(device)

    @classmethod
    def from_pretrained(
        cls,
        model_id_or_path: str,
        *,
        filename: str = "pytorch_model.bin",
        device: str | torch.device = "cpu",
        backbone: Optional[nn.Module] = None,
    ) -> "DinoGeoRegressor":
        """Load weights from a local ``.pth``/``.bin`` file or a HuggingFace repo id.

        The checkpoint is a full state_dict with ``encoder.backbone.*`` and
        ``head.*`` keys (i.e. it includes the frozen backbone weights).
        """
        if os.path.isfile(model_id_or_path):
            weights_path = model_id_or_path
        else:
            from huggingface_hub import hf_hub_download

            weights_path = hf_hub_download(repo_id=model_id_or_path, filename=filename)

        if backbone is None:
            backbone = cls.build_backbone(device)

        model = cls(backbone).to(device)

        if weights_path.endswith(".safetensors"):
            from safetensors.torch import load_file

            state_dict = load_file(weights_path, device=str(device))
        else:
            state_dict = torch.load(weights_path, map_location=device)

        model.load_state_dict(state_dict, strict=True)
        model.eval()
        return model