Create app.py

app.py

@@ -0,0 +1,87 @@
+import gradio as gr
+import torch, numpy as np, json
+from PIL import Image
+from transformers import CLIPProcessor, CLIPModel
+import pygeohash as pgh
+import os
+
+EXPORT_DIR = "."
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TOP_K = 3
+
+# ---------------- Load metadata ----------------
+metadata = json.load(open(os.path.join(EXPORT_DIR, "metadata.json")))
+geoh2id = metadata["geoh2id"]
+id2geoh = {int(v): k for k,v in geoh2id.items()}
+country2id = metadata["country2id"]
+clip_model_name = metadata["clip_model"]
+dim = metadata["embedding_dim"]
+num_fine = len(geoh2id)
+num_countries = len(country2id)
+
+# ---------------- Model definition ----------------
+import torch.nn as nn
+class GeoHybridModel(nn.Module):
+    def __init__(self, in_dim, num_classes, num_countries, hidden=1024, drop=0.3):
+        super().__init__()
+        self.shared = nn.Sequential(
+            nn.Linear(in_dim, hidden), nn.ReLU(), nn.Dropout(drop),
+            nn.Linear(hidden, hidden//2), nn.ReLU(), nn.Dropout(drop)
+        )
+        self.classifier = nn.Linear(hidden//2, num_classes)
+        self.regressor  = nn.Linear(hidden//2, 2)
+        self.country_classifier = nn.Linear(hidden//2, num_countries)
+
+    def forward(self, x):
+        feat = self.shared(x)
+        return self.classifier(feat), self.regressor(feat), self.country_classifier(feat)
+
+# Load weights
+model = GeoHybridModel(dim, num_fine, num_countries)
+model.load_state_dict(torch.load(os.path.join(EXPORT_DIR, "model.pt"), map_location=DEVICE))
+model.to(DEVICE).eval()
+
+# Load CLIP
+clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
+clip_model = CLIPModel.from_pretrained(clip_model_name).to(DEVICE).eval()
+
+# ---------------- Haversine ----------------
+def haversine(lat1, lon1, lat2, lon2):
+    R = 6371.0
+    phi1,phi2 = np.radians(lat1), np.radians(lat2)
+    dphi = np.radians(lat2-lat1)
+    dlambda = np.radians(lon2-lon1)
+    a = np.sin(dphi/2)**2 + np.cos(phi1)*np.cos(phi2)*np.sin(dlambda/2)**2
+    return 2*R*np.arctan2(np.sqrt(a), np.sqrt(1-a))
+
+# ---------------- Prediction ----------------
+def predict_geohash(img: Image.Image):
+    c_in = clip_processor(images=img, return_tensors="pt").to(DEVICE)
+    with torch.no_grad():
+        emb = clip_model.get_image_features(**c_in)
+        emb = emb / emb.norm(p=2, dim=-1, keepdim=True)
+        out_class, out_offset, _ = model(emb)
+        out_class_np = out_class.cpu().numpy()[0]
+        out_offset_np = out_offset.cpu().numpy()[0]
+
+    topk_idx = out_class_np.argsort()[-TOP_K:][::-1]
+    preds = []
+    for i in topk_idx:
+        geoh = id2geoh[i]
+        lat_base, lon_base, cell_lat, cell_lon = pgh.decode_exactly(geoh)
+        lat_pred = lat_base + out_offset_np[0]*cell_lat
+        lon_pred = lon_base + out_offset_np[1]*cell_lon
+        preds.append(f"{geoh} → {lat_pred:.5f},{lon_pred:.5f}")
+    return "\n".join(preds)
+
+# ---------------- Gradio UI ----------------
+iface = gr.Interface(
+    fn=predict_geohash,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Textbox(),
+    title="Locus - GeoGuessr Image to Coordinates model",
+    description="Upload a streetview image and get top-K predicted geohashes with lat/lon."
+)
+
+if __name__ == "__main__":
+    iface.launch()