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()}
 
@@ -16,11 +17,8 @@ transform = transforms.Compose([
 ])
 
 # Function to load and preprocess an image
-def load_image(image):
-    if isinstance(image, str):  # If image is a file path
-        image = Image.open(image).convert('RGB')
-    elif isinstance(image, np.ndarray):  # If image is a numpy array
-        image = Image.fromarray(image.astype('uint8'), 'RGB')
+def load_image(image_path):
+    image = Image.open(image_path).convert('RGB')
     image = transform(image).unsqueeze(0)  # Add batch dimension
     return image
 
@@ -33,12 +31,9 @@ model.load_state_dict(torch.load('ARTmodelo5ns_vit_weights_epoch6.pth', map_loca
 model.eval()
 
 # Define the prediction function
-def predict_image(image):
-    if image is None:
-        return "No image provided", None
-    
-    # Load and transform the image
-    img_tensor = load_image(image)
+def predict_image(image_path):
+    # Load and transform the image from the file path
+    img_tensor = load_image(image_path)
     
     # Perform the prediction
     with torch.no_grad():
@@ -49,9 +44,6 @@ def predict_image(image):
     percentages = probabilities * 100
     results = {idx_to_class[i]: percentages[i].item() for i in range(len(idx_to_class))}
     
-    # Get the highest prediction
-    predicted_label = max(results, key=results.get)
-    
     # Plotting the results
     labels = list(results.keys())
     values = list(results.values())
@@ -65,21 +57,18 @@ def predict_image(image):
     plt.savefig('result.png')
     plt.close(fig)
 
-    return predicted_label, '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="pil", tool="editor", label="Upload an image or take a photo"),
-        gr.Image(source="webcam", type="pil", tool="editor", label="Take a photo")
-    ],
-    outputs=[
-        gr.Textbox(label="Prediction"), 
-        gr.Image(label="Prediction Probabilities")
-    ],
+    inputs=gr.Image(type="filepath", label="Upload an image"),
+    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 with confidence percentages. This model can classify images of skin lesions into one of the following categories: Measles, Chickenpox, Herpes, Melanoma, and Monkeypox. Check out the dataset and paper at: [Link to Dataset](#), [Link to Paper](#)",
+    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
 )