Update app.py

app.py

@@ -3,10 +3,7 @@ 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"
@@ -48,16 +45,7 @@ model.to(DEVICE).eval()
 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 + map ----------------
+# ---------------- Prediction ----------------
 def predict_geohash_map(img: Image.Image):
     c_in = clip_processor(images=img, return_tensors="pt").to(DEVICE)
     with torch.no_grad():
@@ -67,39 +55,30 @@ 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 = []
-    map_center = None
-    fmap = folium.Map(tiles="OpenStreetMap", zoom_start=2)
-    
+    coords = []
+
     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
-        if map_center is None:
-            map_center = [lat_pred, lon_pred]
-            fmap.location = map_center
-            fmap.zoom_start = 6
+        lat_pred = float(lat_base + out_offset_np[0]*cell_lat)
+        lon_pred = float(lon_base + out_offset_np[1]*cell_lon)
+
         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
+        coords.append([lat_pred, lon_pred])
+
+    return "\n".join(preds_text), coords
 
 # ---------------- Gradio UI ----------------
 iface = gr.Interface(
     fn=predict_geohash_map,
     inputs=gr.Image(type="pil"),
-    outputs=[gr.Textbox(label="Top-K Geohash Predictions"), gr.HTML(label="Map")],
+    outputs=[
+        gr.Textbox(label="Top-K Geohash Predictions"),
+        gr.Map(label="Predicted Locations")   # <-- Gradio native map
+    ],
     title="GeoGuessr CLIP Top-K Predictor",
     description="Upload a streetview image and see top-K predicted geohashes and map locations."
 )