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