| """ |
| 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) |
| p_ll = local_xy_to_lonlat(p_m).cpu().tolist() |
|
|
| 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() |
|
|
