improvements

app.py

@@ -6,7 +6,6 @@ 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()}
 
@@ -31,9 +30,16 @@ model.load_state_dict(torch.load('ARTmodelo5ns_vit_weights_epoch6.pth', map_loca
 model.eval()
 
 # Define the prediction function
-def predict_image(image_path):
+def predict_image(image):
+    if image is None:
+        return "No image provided", None
+    
+    # Convert the image to PIL if it's not a filepath
+    if isinstance(image, np.ndarray):
+        image = Image.fromarray(image.astype('uint8'), 'RGB')
+    
     # Load and transform the image from the file path
-    img_tensor = load_image(image_path)
+    img_tensor = transform(image).unsqueeze(0)
     
     # Perform the prediction
     with torch.no_grad():
@@ -44,6 +50,9 @@ def predict_image(image_path):
     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())
@@ -57,18 +66,21 @@ def predict_image(image_path):
     plt.savefig('result.png')
     plt.close(fig)
 
-    # Return the top prediction as well
-    top_prediction = max(results, key=results.get)
-    
-    return top_prediction, 'result.png'
+    return predicted_label, 'result.png'
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict_image,
-    inputs=gr.Image(type="filepath", label="Upload an image"),
-    outputs=[gr.Textbox(label="Prediction"), gr.Image(label="Prediction Probabilities")],
+    inputs=[
+        gr.Image(source="upload", 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")
+    ],
     title="Skin Lesion Image Classification",
-    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](#)",
+    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](#)",
     theme="huggingface",
     live=True
 )