File size: 8,774 Bytes
2242b6a | 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 194 195 196 197 198 199 200 201 202 203 204 205 | """
High-level inference for Img2LocJakarta 2.0 (DINOv3 + conditional flow matching).
Example
-------
from inference import GeoTagPredictor
predictor = GeoTagPredictor("nadh0708/Img2LocJakarta2.0") # HF repo id
# ...or a local folder / .safetensors file:
predictor = GeoTagPredictor(".")
coords = predictor.predict("street.jpg", num_particles=100)
# 100 predicted coordinates: [{'lat': -6.22, 'lon': 106.83}, ...]
clouds = predictor.predict(["a.jpg", "b.jpg"]) # one list of coords per image
predictor.summarize(coords) # mean point + sigma / p50 / p90
geojson = predictor.to_geojson(["a.jpg"], num_particles=100) # 100 Point features
The model is generative: `num_particles` Gaussian samples are flowed onto the target and
the arrival points are the predicted coordinates (N particles -> N coordinates). Predictions are only
meaningful for Jakarta street imagery.
CLI
---
python inference.py nadh0708/Img2LocJakarta2.0 street.jpg -o out.geojson
"""
from __future__ import annotations
import io
import math
from typing import Optional, Union
import torch
from PIL import Image
from torchvision.transforms import v2
from modeling_geotag import DinoV3CFMRegressor, local_xy_to_lonlat
_IMAGENET_MEAN = (0.485, 0.456, 0.406)
_IMAGENET_STD = (0.229, 0.224, 0.225)
INPUT_SIZE = 256
ImageInput = Union[str, bytes, Image.Image]
class GeoTagPredictor:
def __init__(
self,
model_id_or_path: str,
device: str | torch.device | None = None,
filename: str = "model.safetensors",
dinov3_dir: Optional[str] = None,
):
self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
self.model = DinoV3CFMRegressor.from_pretrained(
model_id_or_path, filename=filename, device=self.device, dinov3_dir=dinov3_dir
)
self.transform = v2.Compose([
v2.ToImage(),
v2.Resize((INPUT_SIZE, INPUT_SIZE), antialias=True),
v2.ToDtype(torch.float32, scale=True),
v2.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
])
@staticmethod
def _load(image: ImageInput) -> Image.Image:
if isinstance(image, Image.Image):
return image.convert("RGB")
if isinstance(image, (bytes, bytearray)):
return Image.open(io.BytesIO(image)).convert("RGB")
if isinstance(image, str) and image.startswith(("http://", "https://")):
import urllib.request
req = urllib.request.Request(image, headers={"User-Agent": "img2loc"})
with urllib.request.urlopen(req, timeout=30) as r:
return Image.open(io.BytesIO(r.read())).convert("RGB")
return Image.open(image).convert("RGB")
@torch.no_grad()
def predict(
self,
images: Union[ImageInput, list[ImageInput]],
num_particles: int = 256,
steps: int = 8,
seed: Optional[int] = 0,
) -> Union[list[dict], list[list[dict]]]:
"""Predicted coordinates: `num_particles` samples of the location distribution.
One image -> list of `num_particles` dicts ``{"lat", "lon"}``.
List input -> one such list per image.
Use `summarize()` for the mean point and uncertainty of a cloud.
"""
single = not isinstance(images, (list, tuple))
batch = [images] if single else list(images)
x = torch.stack([self.transform(self._load(im)) for im in batch]).to(self.device)
gen = None
if seed is not None:
gen = torch.Generator(device=self.device).manual_seed(seed)
p_m = self.model.sample_particles(x, num_particles, steps, gen) # (B, N, 2) metres
p_ll = local_xy_to_lonlat(p_m).cpu().tolist() # (B, N, 2) [lon, lat]
out = [[{"lat": lat, "lon": lon} for lon, lat in cloud] for cloud in p_ll]
return out[0] if single else out
@staticmethod
def summarize(coords: list[dict]) -> dict:
"""Mean point + uncertainty of one cloud returned by `predict`: lat/lon of the mean,
sigma_m (RMS spread), p50/p90_radius_m (radius around the mean holding 50/90 % of
the coordinates), n_particles."""
import numpy as np
lat = np.array([c["lat"] for c in coords])
lon = np.array([c["lon"] for c in coords])
m_lat, m_lon = lat.mean(), lon.mean()
dx = (lon - m_lon) * 111320 * math.cos(math.radians(m_lat))
dy = (lat - m_lat) * 110574
d = np.hypot(dx, dy)
return {
"lat": float(m_lat), "lon": float(m_lon), "n_particles": len(coords),
"sigma_m": float(np.sqrt((d ** 2).sum() / max(1, len(d) - 1))),
"p50_radius_m": float(np.percentile(d, 50)),
"p90_radius_m": float(np.percentile(d, 90)),
}
def to_geojson(
self,
images: Union[ImageInput, list[ImageInput]],
names: Optional[list[str]] = None,
truth: Optional[tuple[float, float]] = None,
**kwargs,
) -> dict:
"""FeatureCollection with one Point per predicted coordinate (so 100 particles ->
100 features per image). Each feature carries `image` and `particle` (index);
the top-level `summary` member lists mean/sigma/radii per image.
`truth=(lon, lat)` adds `error_m` to every feature and to the summary."""
batch = [images] if not isinstance(images, (list, tuple)) else list(images)
names = names or [im if isinstance(im, str) else f"image_{i}" for i, im in enumerate(batch)]
clouds = self.predict(batch, **kwargs)
def err(lon, lat):
dx = (lon - truth[0]) * 111320 * math.cos(math.radians(truth[1]))
dy = (lat - truth[1]) * 110574
return math.hypot(dx, dy)
feats, summary = [], []
for name, cloud in zip(names, clouds):
for i, c in enumerate(cloud):
props = {"image": name, "particle": i}
if truth is not None:
props["error_m"] = err(c["lon"], c["lat"])
feats.append({"type": "Feature", "properties": props,
"geometry": {"type": "Point", "coordinates": [c["lon"], c["lat"]]}})
s = {"image": name, **self.summarize(cloud)}
if truth is not None:
s["error_m"] = err(s["lon"], s["lat"])
summary.append(s)
return {"type": "FeatureCollection", "features": feats, "summary": summary}
def _cli() -> None:
import argparse
import json
import sys
parser = argparse.ArgumentParser(description="Geolocate Jakarta street imagery (GeoJSON output).")
parser.add_argument("model", help="HF repo id, local folder or .safetensors path")
parser.add_argument("images", nargs="+", help="image file path(s) or URL(s)")
parser.add_argument("--device", default=None)
parser.add_argument("--filename", default="model.safetensors", help="weights filename inside the HF repo")
parser.add_argument("--dinov3-dir", default=None, help="local clone of facebookresearch/dinov3")
parser.add_argument("--particles", type=int, default=256, help="samples per image")
parser.add_argument("--steps", type=int, default=8, help="Heun integration steps")
parser.add_argument("--seed", type=int, default=0, help="-1 for non-deterministic")
parser.add_argument("--truth", type=float, nargs=2, metavar=("LON", "LAT"),
help="true location; adds error_m to the mean feature")
parser.add_argument("-o", "--output", default=None, help="write GeoJSON here (default: stdout)")
parser.add_argument("--table", action="store_true", help="print a text summary instead of GeoJSON")
args = parser.parse_args()
predictor = GeoTagPredictor(args.model, device=args.device, filename=args.filename,
dinov3_dir=args.dinov3_dir)
kw = dict(num_particles=args.particles, steps=args.steps,
seed=None if args.seed < 0 else args.seed)
if args.table:
for name, cloud in zip(args.images, predictor.predict(args.images, **kw)):
r = predictor.summarize(cloud)
print(f"{name}: {r['n_particles']} coords, mean lon={r['lon']:.5f} lat={r['lat']:.5f} "
f"sigma={r['sigma_m']:.0f}m p50={r['p50_radius_m']:.0f}m p90={r['p90_radius_m']:.0f}m")
return
text = json.dumps(predictor.to_geojson(args.images, truth=args.truth, **kw))
if args.output:
with open(args.output, "w", encoding="utf-8") as f:
f.write(text)
print(f"wrote {args.output}", file=sys.stderr)
else:
print(text)
if __name__ == "__main__":
_cli()
|