Update app.py

app.py

@@ -9,35 +9,21 @@ classification_model = pipeline("text-classification", model="models/text_classi
 masking_model = pipeline("fill-mask", model="models/fill_mask_model", tokenizer="models/fill_mask_model", top_k=100)
 
 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}"
-    response = requests.get(floraveg_url)
-    if response.status_code == 200:
-        soup = BeautifulSoup(response.text, 'html.parser')
-        img_tag = soup.find('img', src=lambda x: x and x.startswith("https://files.ibot.cas.cz/cevs/images/syntaxa/thumbs/"))
-        if img_tag:
-            image_url = img_tag['src']
-        else:
-            image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/a/ac/No_image_available.svg/2048px-No_image_available.svg.png"
-    else:
-        image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/a/ac/No_image_available.svg/2048px-No_image_available.svg.png"
-    image = gr.Image(value=image_url)
-    return image
 
-def return_species_image(species):
-    species = species.capitalize()
-    floraveg_url = f"https://floraveg.eu/taxon/overview/{species}"
+def return_image(task, label):
+    image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/a/ac/No_image_available.svg/2048px-No_image_available.svg.png"
+    if task == "classification":
+        floraveg_url = f"https://floraveg.eu/habitat/overview/{label}"
+        floraveg_tag = "https://files.ibot.cas.cz/cevs/images/syntaxa/thumbs/"
+    elif task == "masking":
+        floraveg_url = f"https://floraveg.eu/taxon/overview/{label.capitalize()}"
+        floraveg_tag = "https://files.ibot.cas.cz/cevs/images/taxa/large/"
     response = requests.get(floraveg_url)
     if response.status_code == 200:
         soup = BeautifulSoup(response.text, 'html.parser')
-        img_tag = soup.find('img', src=lambda x: x and x.startswith("https://files.ibot.cas.cz/cevs/images/taxa/large/"))
+        img_tag = soup.find('img', src=lambda x: x and x.startswith(floraveg_tag))
         if img_tag:
             image_url = img_tag['src']
-        else:
-            image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/a/ac/No_image_available.svg/2048px-No_image_available.svg.png"
-    else:
-        image_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/a/ac/No_image_available.svg/2048px-No_image_available.svg.png"
     image = gr.Image(value=image_url)
     return image
 
@@ -71,7 +57,7 @@ def classification(text, k):
         text = f"This vegetation plot probably belongs to the habitat type {', '.join(habitat_labels[:-1])}, or {habitat_labels[-1]}."
         text += f"\nThe most likely habitat type (i.e., {habitat_labels[0]}) is named '{habitat_name}'."
     text += f"\nSee an image of this habitat type below."
-    image_output = return_habitat_image(habitat_labels[0])
+    image_output = return_image("classification", habitat_labels[0])
     return text, image_output
 
 def masking(text, k):
@@ -125,7 +111,7 @@ def masking(text, k):
         text = f"The most likely missing species are {', '.join(best_predictions[:-1].capitalize())}, and {best_predictions[-1].capitalize()} (positions {', '.join(map(str, best_positions[:-1]))}, and {best_positions[-1]})."
     text += f"\nThe completed vegetation plot is thus '{best_sentence}'."
     text += f"\nSee an image of this species (i.e., {best_predictions[0].capitalize()}) below."
-    image = return_species_image(best_predictions[0])
+    image = return_image("masking", best_predictions[0])
     return text, image
 
 with gr.Blocks() as demo: