Update app.py

app.py

@@ -8,6 +8,19 @@ nlp = spacy.load("./models/en_core_web_sm")
 st.title("SpaGAN Demo")
 st.write("Enter a text, and the system will highlight the geo-entities within it.")
 
+# Define a color map for different entity types
+COLOR_MAP = {
+    'FAC': 'red',    # Facilities (e.g., buildings, airports)
+    'ORG': 'blue',   # Organizations (e.g., companies, institutions)
+    'LOC': 'purple', # Locations (e.g., mountain ranges, water bodies)
+    'GPE': 'green'   # Geopolitical Entities (e.g., countries, cities)
+}
+
+# Display the color key
+st.write("**Color Key:**")
+for label, color in COLOR_MAP.items():
+    st.markdown(f"- **{label}**: <span style='color:{color}'>{color}</span>", unsafe_allow_html=True)
+
 user_input = st.text_area("Input Text", height=200)
 
 # Process the text when the button is clicked
@@ -16,19 +29,18 @@ if st.button("Highlight Geo-Entities"):
         # Process the text using spaCy
         doc = nlp(user_input)
 
-        # Highlight geo-entities
+        # Highlight geo-entities with different colors
         highlighted_text = user_input
         for ent in reversed(doc.ents):
-            if ent.label_ in ['FAC', 'ORG', 'LOC', 'GPE']:  # GPE = Geopolitical Entity, LOC = Location
+            if ent.label_ in COLOR_MAP:
+                color = COLOR_MAP[ent.label_]
                 highlighted_text = (
                     highlighted_text[:ent.start_char] +
-                    f"**:green[{ent.text}]**" + 
+                    f"<span style='color:{color}; font-weight:bold'>{ent.text}</span>" + 
                     highlighted_text[ent.end_char:]
                 )
 
-        # Display the highlighted text
-        st.markdown(highlighted_text)
+        # Display the highlighted text with HTML support
+        st.markdown(highlighted_text, unsafe_allow_html=True)
     else:
-        st.error("Please enter some text.")
-
-st.write("Note: The model identifies and highlights geo-entities.")
+        st.error("Please enter some text.")