Update app.py

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
-import torch, numpy as np, json
+import torch, numpy as np, json, os
 from PIL import Image
 from transformers import CLIPProcessor, CLIPModel
 import pygeohash as pgh
-import os
 import plotly.graph_objects as go
+import torch.nn as nn
 
 EXPORT_DIR = "."
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -20,8 +20,7 @@ dim = metadata["embedding_dim"]
 num_fine = len(geoh2id)
 num_countries = len(country2id)
 
-# ---------------- Model definition ----------------
-import torch.nn as nn
+# ---------------- Model ----------------
 class GeoHybridModel(nn.Module):
     def __init__(self, in_dim, num_classes, num_countries, hidden=1024, drop=0.3):
         super().__init__()
@@ -37,12 +36,11 @@ class GeoHybridModel(nn.Module):
         feat = self.shared(x)
         return self.classifier(feat), self.regressor(feat), self.country_classifier(feat)
 
-# Load model 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 ----------------
 clip_processor = CLIPProcessor.from_pretrained(clip_model_name)
 clip_model = CLIPModel.from_pretrained(clip_model_name).to(DEVICE).eval()
 
@@ -55,8 +53,8 @@ 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))
 
-# ---------------- Prediction + Plotly map ----------------
-def predict_geohash_map(img: Image.Image):
+# ---------------- Prediction + Map ----------------
+def predict_geohash_map(img: Image.Image, true_lat=None, true_lon=None):
     c_in = clip_processor(images=img, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
         emb = clip_model.get_image_features(**c_in)
@@ -74,42 +72,68 @@ def predict_geohash_map(img: Image.Image):
         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}")
         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",
+    # Base map
+    fig = go.Figure()
+
+    # Prediction markers
+    fig.add_trace(go.Scattermapbox(
+        lat=lats, lon=lons, mode="markers+text",
+        text=labels, textposition="top right",
         marker=go.scattermapbox.Marker(size=12, color="blue"),
+        name="Predictions"
     ))
 
+    # True location + lines
+    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"
+        ))
+        # 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)
+            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",
+                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]
+    center_lon = true_lon if true_lon is not None else lons[0]
     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
-        ),
+        mapbox=dict(center=go.layout.mapbox.Center(lat=center_lat, lon=center_lon), zoom=4),
         margin={"r":0,"t":0,"l":0,"b":0}
     )
 
     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.Plot(label="Map")],
-    title="GeoGuessr CLIP Top-K Predictor",
-    description="Upload a streetview image and see top-K predicted geohashes and map locations."
-)
+with gr.Blocks() as demo:
+    with gr.Row():
+        img_in = gr.Image(type="pil", label="Upload Streetview Image")
+        with gr.Column():
+            lat_in = gr.Number(label="True Latitude (optional)")
+            lon_in = gr.Number(label="True Longitude (optional)")
+    out_text = gr.Textbox(label="Predictions + Distances")
+    out_plot = gr.Plot(label="Map")
+
+    run_btn = gr.Button("Run Prediction")
+    run_btn.click(predict_geohash_map, [img_in, lat_in, lon_in], [out_text, out_plot])
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()