Update app.py

app.py

@@ -4,6 +4,7 @@ from PIL import Image
 from transformers import CLIPProcessor, CLIPModel
 import pygeohash as pgh
 import os
+import plotly.graph_objects as go
 
 EXPORT_DIR = "."
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -45,7 +46,16 @@ model.to(DEVICE).eval()
 clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
 clip_model = CLIPModel.from_pretrained(clip_model_name).to(DEVICE).eval()
 
-# ---------------- Prediction ----------------
+# ---------------- 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 + Plotly map ----------------
 def predict_geohash_map(img: Image.Image):
     c_in = clip_processor(images=img, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
@@ -55,30 +65,48 @@ def predict_geohash_map(img: Image.Image):
         out_class_np = out_class.cpu().numpy()[0]
         out_offset_np = out_offset.cpu().numpy()[0]
 
-    # Get top-k class predictions
     topk_idx = out_class_np.argsort()[-TOP_K:][::-1]
     preds_text = []
-    coords = []
+    lats, lons, labels = [], [], []
 
     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 = float(lat_base + out_offset_np[0]*cell_lat)
-        lon_pred = float(lon_base + out_offset_np[1]*cell_lon)
+        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}")
-        coords.append([lat_pred, lon_pred])
+        lats.append(lat_pred)
+        lons.append(lon_pred)
+        labels.append(f"Top-{rank}: {geoh}")
+
+    # Plotly scattermapbox
+    fig = go.Figure(go.Scattermapbox(
+        lat=lats,
+        lon=lons,
+        mode="markers+text",
+        text=labels,
+        textposition="top right",
+        marker=go.scattermapbox.Marker(size=12, color="blue"),
+    ))
+
+    fig.update_layout(
+        mapbox_style="open-street-map",
+        hovermode="closest",
+        mapbox=dict(
+            center=go.layout.mapbox.Center(lat=lats[0], lon=lons[0]),
+            zoom=4
+        ),
+        margin={"r":0,"t":0,"l":0,"b":0}
+    )
 
-    return "\n".join(preds_text), coords
+    return "\n".join(preds_text), fig
 
 # ---------------- Gradio UI ----------------
 iface = gr.Interface(
     fn=predict_geohash_map,
     inputs=gr.Image(type="pil"),
-    outputs=[
-        gr.Textbox(label="Top-K Geohash Predictions"),
-        gr.Map(label="Predicted Locations")   # <-- Gradio native map
-    ],
+    outputs=[gr.Textbox(label="Top-K Geohash Predictions"), gr.Plot(label="Map")],
     title="GeoGuessr CLIP Top-K Predictor",
     description="Upload a streetview image and see top-K predicted geohashes and map locations."
 )