File size: 6,831 Bytes
1bb5a3e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | """
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
|