Update app.py

app.py

@@ -5,15 +5,16 @@ from transformers import CLIPProcessor, CLIPModel
 import pygeohash as pgh
 import plotly.graph_objects as go
 import torch.nn as nn
+import reverse_geocoder as rg
+import pycountry  # for full country names
 
 EXPORT_DIR = "."
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-TOP_K = 5
 
 # ---------------- 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()}
+id2geoh = {int(v): k for k, v in geoh2id.items()}
 country2id = metadata["country2id"]
 clip_model_name = metadata["clip_model"]
 dim = metadata["embedding_dim"]
@@ -53,8 +54,18 @@ def haversine(lat1, lon1, lat2, lon2):
     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))
 
+# ---------------- Country lookup ----------------
+def latlon_to_country(lat, lon):
+    try:
+        res = rg.search((lat, lon))[0]
+        cc = res['cc']
+        country = pycountry.countries.get(alpha_2=cc)
+        return country.name if country else cc, cc
+    except:
+        return "Unknown", "??"
+
 # ---------------- Prediction + Map ----------------
-def predict_geohash_map(img: Image.Image, true_lat=None, true_lon=None):
+def predict_geohash_map(img: Image.Image, true_lat=None, true_lon=None, top_k=5):
     c_in = clip_processor(images=img, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
         emb = clip_model.get_image_features(**c_in)
@@ -63,52 +74,67 @@ def predict_geohash_map(img: Image.Image, true_lat=None, true_lon=None):
         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_text = []
-    lats, lons, labels = [], [], []
+    topk_idx = out_class_np.argsort()[-top_k:][::-1]
+    lats, lons = [], []
+
+    # Markdown table header
+    preds_text = ["| Rank | Country | Code | Latitude | Longitude | Distance |",
+                  "|------|---------|------|----------|-----------|----------|"]
 
     for rank, i in enumerate(topk_idx, 1):
         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_text.append(f"{rank}. {geoh} → {lat_pred:.5f},{lon_pred:.5f}")
+        lat_pred = lat_base + out_offset_np[0] * cell_lat
+        lon_pred = lon_base + out_offset_np[1] * cell_lon
+
+        country_name, country_code = latlon_to_country(lat_pred, lon_pred)
+
+        dist_str = ""
+        if true_lat is not None and true_lon is not None:
+            dist = haversine(true_lat, true_lon, lat_pred, lon_pred)
+            if dist < 100:
+                dist_str = f"🟢 {dist:.1f} km"
+            elif dist < 1000:
+                dist_str = f"🟡 {dist:.1f} km"
+            else:
+                dist_str = f"🔴 {dist:.1f} km"
+
+        preds_text.append(
+            f"| {rank} | {country_name} | {country_code} | {lat_pred:.5f} | {lon_pred:.5f} | {dist_str} |"
+        )
+
         lats.append(lat_pred)
         lons.append(lon_pred)
-        labels.append(f"Top-{rank}: {geoh}")
 
-    # Base map
+    # -------- Map Plot --------
     fig = go.Figure()
 
-    # Prediction markers
+    # Predictions
     fig.add_trace(go.Scattermapbox(
         lat=lats, lon=lons, mode="markers+text",
-        text=labels, textposition="top right",
+        text=[f"Top-{i+1}" for i in range(len(lats))],
+        textposition="top right",
         marker=go.scattermapbox.Marker(size=12, color="blue"),
-        name="Predictions"
+        name="Predictions", showlegend=False
     ))
 
-    # True location + lines
+    # True location (optional)
     if true_lat is not None and true_lon is not None:
         fig.add_trace(go.Scattermapbox(
             lat=[true_lat], lon=[true_lon], mode="markers+text",
             text=["True Location"], textposition="bottom right",
             marker=go.scattermapbox.Marker(size=14, color="red"),
-            name="True Location"
+            name="True Location", showlegend=False
         ))
-        # Add connecting lines + distance annotations
-        for rank, (lat_p, lon_p) in enumerate(zip(lats, lons), 1):
-            dist = haversine(true_lat, true_lon, lat_p, lon_p)
+        # Connect with lines
+        for lat_p, lon_p in zip(lats, lons):
             fig.add_trace(go.Scattermapbox(
                 lat=[true_lat, lat_p],
                 lon=[true_lon, lon_p],
-                mode="lines+text",
-                text=[None, f"{dist:.1f} km"],
-                textposition="middle right",
+                mode="lines",
                 line=dict(width=2, color="green"),
                 showlegend=False
             ))
-            preds_text.append(f"Dist to Top-{rank}: {dist:.1f} km")
 
     # Layout
     center_lat = true_lat if true_lat is not None else lats[0]
@@ -117,7 +143,8 @@ def predict_geohash_map(img: Image.Image, true_lat=None, true_lon=None):
         mapbox_style="open-street-map",
         hovermode="closest",
         mapbox=dict(center=go.layout.mapbox.Center(lat=center_lat, lon=center_lon), zoom=4),
-        margin={"r":0,"t":0,"l":0,"b":0}
+        margin={"r": 0, "t": 0, "l": 0, "b": 0},
+        showlegend=False
     )
 
     return "\n".join(preds_text), fig
@@ -126,27 +153,30 @@ def predict_geohash_map(img: Image.Image, true_lat=None, true_lon=None):
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-        # 🌍 Locus - GeoGuessr Image to Coordinates Model  
+        # 🌍 Locus - GeoGuessr Image to Coordinates Model
         Upload a **streetview-style image** (photo or screenshot).  
         The model can handle images with a small UI overlay (like from Google Street View),  
-        as long as the **main scene is clearly visible**.  
+        as long as the **main scene is clearly visible**.
 
-        You can optionally provide the actual coordinates to compare predictions and calculate distances on a world map.
-        
-        > **Note:**  
-        > Locus was trained solely on publicly accessible Google Street View images, not on **actual GeoGuessr data**.
+        You can optionally provide the actual coordinates to compare predictions and calculate distances on a world map.  
         """
     )
     with gr.Row():
         img_in = gr.Image(type="pil", label="Upload Streetview Image")
         with gr.Column():
-            lat_in = gr.Number(label="True Latitude (optional)", info="Enter the known latitude if available.")
-            lon_in = gr.Number(label="True Longitude (optional)", info="Enter the known longitude if available.")
-    out_text = gr.Textbox(label="Predictions + Distances")
+            lat_in = gr.Number(label="True Latitude (optional)")
+            lon_in = gr.Number(label="True Longitude (optional)")
+            topk_in = gr.Slider(1, 10, value=5, step=1,
+                                label="Top K Predictions",
+                                info="Choose how many top predictions to display (1–10).")
+
+    out_text = gr.Markdown(label="Predictions (Markdown Table)")
     out_plot = gr.Plot(label="Map")
 
     run_btn = gr.Button("Run Prediction", variant="primary")
-    run_btn.click(predict_geohash_map, [img_in, lat_in, lon_in], [out_text, out_plot])
+    run_btn.click(predict_geohash_map,
+                  [img_in, lat_in, lon_in, topk_in],
+                  [out_text, out_plot])
 
 if __name__ == "__main__":
     demo.launch()