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