#!/usr/bin/env python3 """ Chipoint — run the image-geolocation model on your own photo(s). Encodes each image with three frozen vision encoders, projects through trained heads, retrieves the nearest gallery images, and applies the density-weighted consensus that produced the benchmark numbers. Prints a predicted (latitude, longitude) per image. Batch-friendly: pass many images and each 5 GB gallery index is read exactly ONCE for the whole batch (all images are scored against it in a single pass), instead of re-reading 15 GB per image. Usage: python run_chipoint.py path/to/photo.jpg [more.jpg ...] # globs are expanded, e.g. photos/*.jpg Weights (heads + gallery indexes + coords) must sit next to this script (as when you download the repo). Override the location with CHIPOINT_DIR=/path/to/weights. DINOv3 is pulled ungated from ModelScope unless DINOV3_PATH points at a local copy. """ import os, sys, glob os.environ.setdefault("PYTORCH_HIP_ALLOC_CONF", "expandable_segments:True") os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") from pathlib import Path import numpy as np, torch, torch.nn as nn, torch.nn.functional as F from PIL import Image W = Path(os.environ.get("CHIPOINT_DIR") or Path(__file__).resolve().parent) # weights live next to this script DEV = "cuda" if torch.cuda.is_available() else "cpu" AC = torch.bfloat16 if DEV == "cuda" else torch.float32 M, K = 200, 64 # retrieve top-M per encoder, keep top-K fused candidates Gv, LAM, MU, BW1, BW2 = 0.8, 1.7, 0.8, 1492.7, 750.0 # locked density-weighted-consensus config ENC_D = {"so400m256": 1152, "dinov3": 1024, "gopt": 1536} class Head(nn.Module): def __init__(s, D): super().__init__(); s.net = nn.Sequential(nn.Linear(D, 1024), nn.GELU(), nn.Dropout(0.1), nn.Linear(1024, 512)) def forward(s, x): return F.normalize(s.net(x), dim=-1) def load_head(tag, D): h = Head(D); sd = torch.load(W / f"proj_head_{tag}.pt", map_location="cpu") h.load_state_dict({k: v for k, v in sd.items() if k.startswith("net.")}) # drop geocell centers, unused at inference return h.to(DEV).eval() # ---- encode ALL images with the 3 encoders, then release the encoders (free VRAM for retrieval) ---------- @torch.no_grad() def encode_all(imgs, bs=64): import open_clip, torchvision.transforms as T from transformers import AutoModel, AutoImageProcessor E = {} for tag, name in [("so400m256", "ViT-SO400M-16-SigLIP2-256"), ("gopt", "ViT-gopt-16-SigLIP2-256")]: m, _, pre = open_clip.create_model_and_transforms(name, pretrained="webli"); m = m.to(DEV).eval(); out = [] for i in range(0, len(imgs), bs): x = torch.stack([pre(p) for p in imgs[i:i+bs]]).to(DEV) with torch.autocast(DEV, dtype=AC): e = m.encode_image(x).float() out.append(F.normalize(e, dim=-1).cpu().numpy().astype("float32")) E[tag] = np.concatenate(out); del m if DEV == "cuda": torch.cuda.empty_cache() src = os.environ.get("DINOV3_PATH") if not src: from modelscope import snapshot_download as _ms src = _ms("facebook/dinov3-vitl16-pretrain-lvd1689m") dm = AutoModel.from_pretrained(src).to(DEV).eval() if DEV == "cuda": dm = dm.half() proc = AutoImageProcessor.from_pretrained(src) mean = torch.tensor(proc.image_mean).view(1, 3, 1, 1); std = torch.tensor(proc.image_std).view(1, 3, 1, 1) pre = T.Compose([T.Resize(256), T.CenterCrop(224), T.ToTensor()]); out = [] for i in range(0, len(imgs), bs): x = torch.stack([(pre(p) - mean[0]) / std[0] for p in imgs[i:i+bs]]).to(DEV) x = x.half() if DEV == "cuda" else x o = dm(x); e = o.pooler_output if getattr(o, "pooler_output", None) is not None else o.last_hidden_state[:, 0] out.append(F.normalize(e.float(), dim=-1).cpu().numpy().astype("float32")) E["dinov3"] = np.concatenate(out); del dm if DEV == "cuda": torch.cuda.empty_cache() return E @torch.no_grad() def project(tag, emb): return F.normalize(HEADS[tag](torch.tensor(emb).to(DEV)), dim=-1).cpu().numpy().astype("float32") # ---- retrieve ALL queries against one gallery in a SINGLE pass (gallery read once) ----------------------- @torch.no_grad() def retrieve_all(tag, Q): # Q:[N,512] -> idx[N,M], sim[N,M] g = np.load(W / f"proj_head_gallery_512_{tag}.npy", mmap_mode="r"); N = len(Q) if DEV == "cuda": qt = torch.from_numpy(Q).to(DEV).half() bs = torch.full((N, M), -1e9, device=DEV); bi = torch.zeros((N, M), dtype=torch.long, device=DEV); step = 250000 for i in range(0, g.shape[0], step): gg = torch.from_numpy(np.asarray(g[i:i+step], dtype="float16")).to(DEV) s = (qt @ gg.T).float(); cs, ci = s.topk(min(M, s.shape[1]), 1); ci = ci + i al = torch.cat([bs, cs], 1); ai = torch.cat([bi, ci], 1); o = al.topk(M, 1).indices bs = torch.gather(al, 1, o); bi = torch.gather(ai, 1, o); del gg, s torch.cuda.empty_cache() return bi.cpu().numpy(), bs.cpu().numpy() bs = np.full((N, M), -1e9, "float32"); bi = np.zeros((N, M), np.int64); step = 500000 for i in range(0, g.shape[0], step): s = Q @ np.asarray(g[i:i+step], dtype="float32").T ci = np.argpartition(-s, min(M, s.shape[1]-1), 1)[:, :M]; cs = np.take_along_axis(s, ci, 1); ci = ci + i al = np.concatenate([bs, cs], 1); ai = np.concatenate([bi, ci], 1); o = np.argsort(-al, 1)[:, :M] bs = np.take_along_axis(al, o, 1); bi = np.take_along_axis(ai, o, 1) return bi, bs def zc(a): return (a - a.mean()) / (a.std() + 1e-6) def cell_key(ll, C): return (np.round(ll[..., 0]/C).astype(np.int64) << 20) ^ (np.round(ll[..., 1]/C).astype(np.int64) & 0xFFFFF) def predict_one(R, n): # fuse the 3 encoders + density-weighted consensus for image n sc = {} for t in ENC_D: idx, sim = R[t]; z = zc(sim[n]) for j, zz in zip(idx[n], z): sc[int(j)] = sc.get(int(j), 0.0) + float(zz) ci = np.array(sorted(sc, key=lambda k: -sc[k])[:K], np.int64); csim = np.array([sc[i] for i in ci], "float32") cc = GC[ci] DW = np.mean([[-np.log(DENS[C].get(int(k), 1) + 1.0) for k in cell_key(cc, C)] for C in (0.25,0.5,1.0,2.0)], 0) la, lo = np.radians(cc[:, 0]), np.radians(cc[:, 1]) D = 2*6371.0*np.arcsin(np.sqrt(np.clip(np.sin((la[:,None]-la[None,:])/2)**2 + np.cos(la)[:,None]*np.cos(la)[None,:]*np.sin((lo[:,None]-lo[None,:])/2)**2, 0, 1))) vote = np.maximum(csim, 0) * np.exp(Gv * zc(DW)) KV1 = (vote[None,:] * np.exp(-D/BW1)).sum(1); KV2 = (vote[None,:] * np.exp(-D/BW2)).sum(1) score = zc(csim) + 0.3*zc(DW) + LAM*zc(KV1) + MU*zc(KV2) return cc[int(score.argmax())] if __name__ == "__main__": paths = [p for a in sys.argv[1:] for p in glob.glob(a)] or sys.exit("usage: python run_chipoint.py IMAGE ...") print(f"device={DEV} loading heads + coords from {W} ...", flush=True) HEADS = {t: load_head(t, D) for t, D in ENC_D.items()} GC = np.load(W / "gallery_all_C_so400m256.npy").astype("float32") # GPS per gallery entry DENS = {C: dict(zip(*[a.tolist() for a in np.unique(cell_key(GC, C), return_counts=True)])) for C in (0.25,0.5,1.0,2.0)} imgs = [Image.open(p).convert("RGB") for p in paths] print(f"encoding {len(imgs)} image(s) ...", flush=True) E = encode_all(imgs) # encoders freed after this print("retrieving (each gallery read once for the whole batch) ...", flush=True) R = {t: retrieve_all(t, project(t, E[t])) for t in ENC_D} # one 5 GB pass per encoder for n, p in enumerate(paths): lat, lon = predict_one(R, n) print(f"{p}\t{lat:.5f}, {lon:.5f}\thttps://maps.google.com/?q={lat:.5f},{lon:.5f}", flush=True)