app.py

@@ -2,10 +2,14 @@ import gradio as gr
 from transformers import pipeline
 import requests
 from bs4 import BeautifulSoup
+import pandas as pd
 
 # Initialize models
 classification_model = pipeline("text-classification", model="models/text_classification_model", tokenizer="models/text_classification_model", top_k=5)
 mask_model = pipeline("fill-mask", model="models/fill_mask_model", tokenizer="models/fill_mask_model", top_k=100)
+
+# Load data
+eunis_habitats = pd.read_excel('data/eunis_habitats.xlsx')
     
 def return_habitat_image(habitat_label):
     floraveg_url = f"https://floraveg.eu/habitat/overview/{habitat_label}"
@@ -70,6 +74,8 @@ def classification(text, k):
         text = f"This vegetation plot belongs to the habitat {habitat_labels[0]}."
     else:
         text = f"This vegetation plot belongs to the habitats {', '.join(habitat_labels[:-1])} and {habitat_labels[-1]}."
+    habitat_name = eunis_habitats[eunis_habitats['EUNIS 2020 code'] == habitat_labels[0]]['EUNIS-2021 habitat name'].values[0]
+    text += f"\nThe most likely habitat is {habitat_name} (see image below)."
     image_output = return_habitat_image(habitat_labels[0])
     return text, image_output
 
@@ -82,9 +88,7 @@ def masking(text):
     best_position = None
     best_sentence = None
 
-    # Loop through each position in the sentence
     for i in range(len(text_split) + 1):
-        # Create masked text
         masked_text = ', '.join(text_split[:i] + ['[MASK]'] + text_split[i:])
         
         j = 0
@@ -99,7 +103,6 @@ def masking(text):
         score = prediction['score']
         sentence = prediction['sequence']
 
-        # Update best prediction and position if score is higher
         if score > max_score:
             max_score = score
             best_prediction = species
@@ -119,7 +122,7 @@ with gr.Blocks() as demo:
         with gr.Row():
             with gr.Column():
                 species = gr.Textbox(lines=2, label="Species", placeholder="Enter a list of comma-separated binomial names here.")
-                top_k = gr.Slider(1, 5, value=1, label="Top-k", info="Choose between 1 and 5.")
+                top_k = gr.Slider(1, 5, value=1, label="Top-k", info="Choose the number of top habitats to display.")
             with gr.Column():
                 text_output_1 = gr.Textbox()
                 text_output_2 = gr.Image()
@@ -138,7 +141,7 @@ with gr.Blocks() as demo:
         gr.Markdown("""<h5 style="text-align: center;">An example of input</h5>""")
         gr.Examples([["vaccinium myrtillus, dryopteris dilatata, molinia caerulea"]], [species_2], [image_output_1, image_output_2], masking, True)
 
-    text_button.click(classification, inputs=[species], outputs=[text_output_1, text_output_2])
+    text_button.click(classification, inputs=[species, top_k], outputs=[text_output_1, text_output_2])
     image_button.click(masking, inputs=[species_2], outputs=[image_output_1, image_output_2])
 
 demo.launch()