Update app.py

app.py

@@ -7,7 +7,12 @@ import torch.nn as nn
 import copy
 import pydeck as pdk
 import numpy as np
+import base64
 
+def get_base64_encoded_image(image_path):
+    with open(image_path, "rb") as img_file:
+        return base64.b64encode(img_file.read()).decode('utf-8')
+        
 keep_layer_count=6
 
 
@@ -75,7 +80,7 @@ class ByT5ForTextGeotagging(PreTrainedModel):
             return logits
 
 
-@st.cache_data(ttl=None, persist=True)
+@st.cache_resource(ttl=None)
 def load_model_and_tokenizer():
     byt5_tokenizer = AutoTokenizer.from_pretrained("yachay/byt5-geotagging-es", token=st.secrets['token'])
     model = ByT5ForTextGeotagging.from_pretrained("yachay/byt5-geotagging-es", token=st.secrets['token'])
@@ -97,12 +102,21 @@ def geolocate_text_byt5_multiclass(text):
         prob = probas[0][class_idx]
         cumulative_prob += prob
         if cumulative_prob > 0.5:
-            coordinates = model.config.class_to_location.get(str(class_idx))
-            if coordinates:
-                results.append((class_idx, prob, coordinates))
             break
+        coordinates = model.config.class_to_location.get(str(class_idx))
+        if coordinates:
+            results.append((class_idx, prob, coordinates))
+
+    # Check if at least one result is added; if not, add the highest probability class
+    if not results:
+        class_idx = sorted_indices[0]
+        prob = probas[0][class_idx]
+        coordinates = model.config.class_to_location.get(str(class_idx))
+        if coordinates:
+            results.append((class_idx, prob, coordinates))
 
     return results
+
    
 
 def geolocate_text_byt5(text):
@@ -147,37 +161,64 @@ if st.button('Submit'):
 
 if st.session_state.text_input:
     results = geolocate_text_byt5_multiclass(st.session_state.text_input)
+    #st.write(results)
     _, confidence, (lat, lon) = results[0]
-    st.write('Predicted Location: (', lat, lon, '). Confidence: ', 'High' if confidence > 0.2 else 'Low')
+    if len(results) == 1:
+        confidence_def = 'High'
+    elif len(results) < 50:
+        confidence_def = 'Low'
+    else:
+        confidence_def = 'Very low'
+    st.write('Predicted Location: (', lat, lon, '). Confidence: ', confidence_def)
+    if confidence_def == 'Low':
+        st.write('Multiple possible locations were identified as confidence is low')
+    elif confidence_def == 'Very low':
+        st.write('There are too many possible locations as confidence is very low')
 
      # Render map with pydeck
     map_data = pd.DataFrame(
         [[lat, lon]],
         columns=["lat", "lon"]
     )
-
-    layers = [
-        # Your existing layer for the geolocated text
-        pdk.Layer(
-            'ScatterplotLayer',
-            data=map_data,
-            get_position='[lon, lat]',
-            get_color='[200, 30, 0, 160]',
-            get_radius=200,
-        ),
-        # Additional layers for other markers
-    ]
     
-    # Add layers for each additional marker
-    for marker in additional_markers:
-        layer = pdk.Layer(
-            'ScatterplotLayer',
-            data=pd.DataFrame([marker]),
-            get_position='[lon, lat]',
-            get_color=marker['color'],
-            get_radius=200,
-        )
-        layers.append(layer)
+    encoded_image = get_base64_encoded_image("icons8-map-pin-48.png")
+    icon_url = f"data:image/png;base64,{encoded_image}"
+    
+    # Example icon data
+    icon_data = {
+        "url": icon_url,  # URL of the icon image
+        "width": 128,  # Width of the icon in pixels
+        "height": 128,  # Height of the icon in pixels
+        "anchorY": 128  # Anchor point of the icon in pixels (bottom center)
+    }
+
+    # Example location data
+    locations = pd.DataFrame({
+        'lat': [lat],  # Latitude values
+        'lon': [lon],  # Longitude values
+        'icon_data': [icon_data]  # Repeating the icon data for each location
+    })
+
+    layers = []
+        
+    if confidence_def != 'Very low':
+        # Add layers for each additional marker
+        for item in results:
+            _, confidence, (lat, lon) = item
+            layer = pdk.Layer(
+                type='IconLayer',
+                data=pd.DataFrame({
+                    'lat': [lat],  # Latitude values
+                    'lon': [lon],  # Longitude values
+                    'icon_data': [icon_data]  # Repeating the icon data for each location
+                }),
+                get_icon='icon_data',
+                get_size=4,
+                size_scale=15,
+                get_position=['lon', 'lat'],
+                pickable=True
+            )
+            layers.append(layer)
 
     st.pydeck_chart(pdk.Deck(
         map_style='mapbox://styles/mapbox/light-v9',