update percentage breakdown + error message

app.py

@@ -27,25 +27,33 @@ def predict(image):
     label = model.config.id2label[label_index]
 
     # Split the label into leaf category and disease name
-
     label_parts = label.split("___")
     leaf_category = label_parts[0]
 
-    disease_name: str = ""
+    # Assume "background with leaves" as the category and "N/A" as the disease name if label_parts only has one part
     if len(label_parts) == 1:
-        # if label parts only has one part, assume "background with leaves" as the category and "N/A as the disease name"
-        disease_name = "N/A"
+        disease_name = "Not applicable"
     else:
         disease_name = label_parts[1]
 
+
     # Calculate the percentage breakdown of predicted diseases
     percentage_breakdown = {disease: round(float(probability[0, index]) * 100, 2) for index, disease in enumerate(model.config.label2id)}
 
     # Sort the percentage breakdown dictionary by values in descending order
     sorted_percentage = dict(sorted(percentage_breakdown.items(), key=lambda item: item[1], reverse=True))
 
-    # Format the percentage breakdown for printing
-    formatted_percentage = '\n'.join([f"{disease.replace('___', ' Leaf - ').replace('_', ' ').capitalize()}: {percentage}%" for disease, percentage in sorted_percentage.items()])
+    # Get the top 5 predicted diseases
+    top_5_percentage = {disease: percentage for disease, percentage in list(sorted_percentage.items())[:5]}
+
+    # Format the percentage breakdown for printing (top 5 only)
+    formatted_percentage = '\n'.join([f"{disease.replace('___', ' Leaf - ').replace('_', ' ').capitalize()}: {percentage}%" for disease, percentage in top_5_percentage.items()])
+
+    # Include the message in the leaf_category textbox if leaf identification fails
+    if disease_name == "Not applicable":
+        leaf_category = "Unfortunately, we are having trouble to identifying a singular leaf in this image. \n(try uploading a different image with your leaf on a white backgound)"
+        formatted_percentage = "Sorry we could not figure it out :("
+        
 
     return leaf_category.capitalize(), disease_name.replace('_', ' ').capitalize(), formatted_percentage
 
@@ -54,13 +62,12 @@ interface = gr.Interface(
     fn=predict,
     inputs=gr.Image(label="Upload the Image"),
     outputs=[
-        gr.Textbox(label="Leaf Type:"),
+        gr.Textbox(label="Leaf Type:"),  # Modified label to include the message
         gr.Textbox(label="Identified Disease:"),
-        gr.Textbox(label="Percentage Breakdown:")
+        gr.Textbox(label="Percentage Breakdown (top 5):")
     ],
     title="Plant Disease Identifier",
     description="Quick Tip: If the image results do not match, consider taking a picture of the leaf against a clean white background for improved accuracy and better identification of plant diseases."
 )
 
-
 interface.launch(debug=True)  # Start server and launch the UI