Add application file and dependencies

app.py

@@ -4,9 +4,10 @@ import numpy as np
 from torchvision import transforms
 from PIL import Image
 from timm import create_model
+import matplotlib.pyplot as plt
 
-# Define class to index mapping dictionary
-class_to_idx = {'Monkeypox': 0, 'Melanoma': 1, 'Herpes': 2, 'Measles': 3, 'Chickenpox': 4}
+class_to_idx = {'Monkeypox': 0, 'Chickenpox': 1, 'Measles': 2, 'Melanoma': 3, 'Herpes': 4}
+idx_to_class = {v: k for k, v in class_to_idx.items()}
 
 # Data transformation
 transform = transforms.Compose([
@@ -36,21 +37,37 @@ def predict_image(image_path):
     # Perform the prediction
     with torch.no_grad():
         output = model(img_tensor)
-        _, predicted = torch.max(output, 1)
-
-    # Get the predicted label
-    predicted_label = list(class_to_idx.keys())[predicted.item()]
+        probabilities = torch.nn.functional.softmax(output[0], dim=0)
+    
+    # Convert probabilities to percentages
+    percentages = probabilities * 100
+    results = {idx_to_class[i]: percentages[i].item() for i in range(len(idx_to_class))}
     
-    return predicted_label
+    # Plotting the results
+    labels = list(results.keys())
+    values = list(results.values())
+
+    fig, ax = plt.subplots()
+    ax.barh(labels, values, color='skyblue')
+    ax.set_xlabel('Percentage')
+    ax.set_title('Prediction Probabilities')
+
+    plt.tight_layout()
+    plt.savefig('result.png')
+    plt.close(fig)
+
+    return results, '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"),
+    outputs=[gr.Label(label="Prediction"), gr.Image(label="Prediction Probabilities")],
     title="Skin Lesion Image Classification",
-    description="Upload an image of a skin lesion to get a prediction."
+    description="Upload an image of a skin lesion to get a prediction with confidence percentages. Example images: [Link to Example 1](#), [Link to Example 2](#)",
+    theme="huggingface",
+    live=True
 )
 
 # Launch the Gradio interface
-iface.launch()
+iface.launch(share=True)