Update app.py

app.py

@@ -1,20 +1,34 @@
 import streamlit as st
+import spacy
 from transformers import pipeline
 from PIL import Image
 
-pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
+nlp = spacy.load("en_core_web_sm")
 
-st.title("Hot Dog? Or Not?")
+st.title("SpaGAN Demo")
+st.write("Enter a text, and the system will highlight the geo-entities within it.")
 
-file_name = st.file_uploader("Upload a hot dog candidate image")
+user_input = st.text_area("Input Text", height=200)
 
-if file_name is not None:
-    col1, col2 = st.columns(2)
+# Process the text when the button is clicked
+if st.button("Highlight Geo-Entities"):
+    if user_input.strip():
+        # Process the text using spaCy
+        doc = nlp(user_input)
 
-    image = Image.open(file_name)
-    col1.image(image, use_column_width=True)
-    predictions = pipeline(image)
+        # Highlight geo-entities
+        highlighted_text = user_input
+        for ent in reversed(doc.ents):
+            if ent.label_ in ["GPE", "LOC"]:  # GPE = Geopolitical Entity, LOC = Location
+                highlighted_text = (
+                    highlighted_text[:ent.start_char] +
+                    f"**:green[{ent.text}]**" + 
+                    highlighted_text[ent.end_char:]
+                )
 
-    col2.header("Probabilities")
-    for p in predictions:
-        col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")
+        # Display the highlighted text
+        st.markdown(highlighted_text)
+    else:
+        st.error("Please enter some text.")
+
+st.write("Note: The model identifies and highlights geo-entities