Update app.py

app.py

@@ -22,7 +22,7 @@ def return_text(habitat_label, habitat_score, confidence):
         text = f"We can't assign an habitat to this vegetation plot with a confidence of at least {confidence}%."
     return text
     
-def return_image(habitat_label, habitat_score, confidence):
+def return_habitat_image(habitat_label, habitat_score, confidence):
     floraveg_url = f"https://floraveg.eu/habitat/overview/{habitat_label}"
     response = requests.get(floraveg_url)
     if response.status_code == 200:
@@ -39,6 +39,22 @@ def return_image(habitat_label, habitat_score, confidence):
     image = gr.Image(value=image_url)
     return image
 
+def return_species_image(species):
+    species = species[0].capitalize() + species[1:]
+    floraveg_url = f"https://floraveg.eu/taxon/overview/{species}"
+    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/"))
+        if img_tag:
+            image_url = img_tag['src']
+        else:
+            image_url = "https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png"
+    else:
+        image_url = "https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png"
+    image = gr.Image(value=image_url)
+    return image
+
 def classification(text, typology, confidence, task):
     model = return_model(task)
     dataset = return_dataset()
@@ -47,7 +63,7 @@ def classification(text, typology, confidence, task):
     habitat_label = dataset['train'].features['label'].names[int(habitat_label.split('_')[1])]
     habitat_score = result[0]['score']
     formatted_output = return_text(habitat_label, habitat_score, confidence)
-    image_output = return_image(habitat_label, habitat_score, confidence)
+    image_output = return_habitat_image(habitat_label, habitat_score, confidence)
     return formatted_output, image_output
 
 def masking(text, task):
@@ -55,8 +71,9 @@ def masking(text, task):
     masked_text = text + ', [MASK] [MASK]'
     pred = model(masked_text, top_k=1)
     new_species = [pred[i][0]['token_str'] for i in range(len(pred))]
-    text = text + ', ' + ' '.join(new_species)
-    image = gr.Image(value="https://www.salonlfc.com/wp-content/uploads/2018/01/image-not-found-scaled-1150x647.png")
+    new_species = ' '.join(new_species)
+    text = text + ', ' + new_species
+    image = return_species_image(new_species)
     return text, image
 
 def plantbert(text, typology, confidence, task):