improvements

app.py

@@ -6,6 +6,7 @@ from PIL import Image
 from timm import create_model
 import matplotlib.pyplot as plt
 
+# Class to index and index to class mappings
 class_to_idx = {'Monkeypox': 0, 'Measles': 1, 'Chickenpox': 2, 'Herpes': 3, 'Melanoma': 4}
 idx_to_class = {v: k for k, v in class_to_idx.items()}
 
@@ -56,15 +57,18 @@ def predict_image(image_path):
     plt.savefig('result.png')
     plt.close(fig)
 
-    return results, 'result.png'
+    # Return the top prediction as well
+    top_prediction = max(results, key=results.get)
+    
+    return top_prediction, 'result.png'
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="filepath", label="Upload an image"),
-    outputs=[gr.Label(label="Prediction"), gr.Image(label="Prediction Probabilities")],
+    outputs=[gr.Textbox(label="Prediction"), gr.Image(label="Prediction Probabilities")],
     title="Skin Lesion Image Classification",
-    description="Upload an image of a skin lesion to get a prediction . Check out the dataset and paper at: [Link to Example 1](#), [Link to Example 2](#)",
+    description="Upload an image of a skin lesion to get a prediction. This tool helps to classify images of skin lesions into the following categories: Measles, Chickenpox, Herpes, Melanomas, and Monkeypox. Check out the dataset and paper at: [Link to Example 1](#), [Link to Example 2](#)",
     theme="huggingface",
     live=True
 )