add map vis

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 import pandas as pd
 import util
+import folium
 
 # Model names (replace with your actual Hugging Face repo names)
 MODEL_NAMES = {
@@ -28,15 +29,6 @@ def extract_entities(model_choice, text):
 
     if not entities:
         return pd.DataFrame(columns=["outbreak", "cases", "deaths", "date", "location", "latitude", "longtitude"])
-
-    # Format results into DataFrame
-    data = [{
-        "Entity": ent["word"],
-        "Label": ent["entity_group"],
-        "Score": f"{ent['score']:.2f}",
-        "Start": ent["start"],
-        "End": ent["end"]
-    } for ent in entities]
     
     data = {}
     for ent in entities:
@@ -52,9 +44,27 @@ def extract_entities(model_choice, text):
             location_info = util.get_location(ent['word'])
             data['location'] = location_info['name']
             data['latitude'] = location_info['latitude']
-            data['longtitude'] = location_info['longitude'] 
+            data['longtitude'] = location_info['longitude']
+    
+    m = None
+    if 'latitude' in data:
+        m = folium.Map(location=[data['latitude'], data['longtitude']])
+        info = f"""
+        <h3>Outbreak Information</h3>
+        <strong>Location:</strong> {data.get('location', 'N/A')}<br>
+        <strong>Outbreak:</strong> {data.get('outbreak', 'N/A')}<br>
+        <strong>Cases:</strong> {data.get('cases', 'N/A')}<br>
+        <strong>Deaths:</strong> {data.get('deaths', 'N/A')}<br>
+        <strong>Date:</strong> {data.get('date', 'N/A')}<br>
+        <strong>Location:</strong> {data.get('location', 'N/A')}
+        """
+        folium.Marker(
+            location=[data['latitude'], data['longtitude']],
+            popup=info,
+            icon=folium.Icon(icon="warning", color="red"),
+        ).add_to(m)
 
-    return pd.DataFrame([data])
+    return pd.DataFrame([data]), m._repr_html_()
 
 # Sample example texts for testing
 examples = [
@@ -71,17 +81,17 @@ with gr.Blocks() as demo:
     with gr.Row():
         model_dropdown = gr.Dropdown(choices=list(MODEL_NAMES.keys()), label="Select Model", value="BioBERT")
         text_input = gr.Textbox(label="Input Text", placeholder="Enter your sentence here...", lines=4)
-
-    output_table = gr.Dataframe(headers=["outbreak", "cases", "deaths", "date", "location", "latitude", "longtitude"], interactive=False)
-
-    run_button = gr.Button("Extract Entities")
-
-    run_button.click(fn=extract_entities, inputs=[model_dropdown, text_input], outputs=output_table)
-
+        
     gr.Examples(
         examples=examples,
         inputs=[text_input],
         label="Try Examples",
     )
+    
+    run_button = gr.Button("Extract Entities")
+    output_table = gr.Dataframe(headers=["outbreak", "cases", "deaths", "date", "location", "latitude", "longtitude"], interactive=False)
+    output_map = gr.HTML(label="Map")
+
+    run_button.click(fn=extract_entities, inputs=[model_dropdown, text_input], outputs=[output_table, output_map])
 
 demo.launch()