Update app.py

app.py

@@ -3,7 +3,10 @@ import torch, numpy as np, json
 from PIL import Image
 from transformers import CLIPProcessor, CLIPModel
 import pygeohash as pgh
+import folium
 import os
+from io import BytesIO
+import base64
 
 EXPORT_DIR = "."
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -36,7 +39,7 @@ class GeoHybridModel(nn.Module):
         feat = self.shared(x)
         return self.classifier(feat), self.regressor(feat), self.country_classifier(feat)
 
-# Load weights
+# 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()
@@ -54,8 +57,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 ----------------
-def predict_geohash(img: Image.Image):
+# ---------------- Prediction + map ----------------
+def predict_geohash_map(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)
@@ -65,22 +68,40 @@ def predict_geohash(img: Image.Image):
         out_offset_np = out_offset.cpu().numpy()[0]
 
     topk_idx = out_class_np.argsort()[-TOP_K:][::-1]
-    preds = []
-    for i in topk_idx:
+    preds_text = []
+    map_center = None
+    fmap = folium.Map(tiles="OpenStreetMap", zoom_start=2)
+    
+    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.append(f"{geoh} → {lat_pred:.5f},{lon_pred:.5f}")
-    return "\n".join(preds)
+        if map_center is None:
+            map_center = [lat_pred, lon_pred]
+            fmap.location = map_center
+            fmap.zoom_start = 6
+        preds_text.append(f"{rank}. {geoh} → {lat_pred:.5f},{lon_pred:.5f}")
+        folium.Marker(
+            location=[lat_pred, lon_pred],
+            popup=f"Top-{rank}: {geoh}",
+            icon=folium.Icon(color="blue" if rank==1 else "green")
+        ).add_to(fmap)
+    
+    # Convert folium map to HTML iframe
+    fmap_file = BytesIO()
+    fmap.save(fmap_file, close_file=False)
+    fmap_html = fmap_file.getvalue().decode()
+    
+    return "\n".join(preds_text), fmap_html
 
 # ---------------- Gradio UI ----------------
 iface = gr.Interface(
-    fn=predict_geohash,
+    fn=predict_geohash_map,
     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."
+    outputs=[gr.Textbox(label="Top-K Geohash Predictions"), gr.HTML(label="Map")],
+    title="GeoGuessr CLIP Top-K Predictor",
+    description="Upload a streetview image and see top-K predicted geohashes and map locations."
 )
 
 if __name__ == "__main__":