Update app.py

app.py

@@ -1,36 +1,44 @@
 import gradio as gr
-from PIL import Image
-from model import load_model
-from torchvision import transforms
 import torch
+from torchvision import transforms
+from PIL import Image
+from model import CustomModel, load_model
 
-# Load your model
-model = load_model('best_model_augmented.pth')
-
-# Define the transformations for input images
+# Define the transformation
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
 ])
 
-# Define the class labels
+# Load the model
+model = load_model('best_model_augmented.pth')
+
+# Define the class names
 class_names = ['Normal', 'Monkeypox', 'Chickenpox', 'Measles']
 
 def predict(image):
     # Preprocess the image
-    image = transform(image).unsqueeze(0)  # Add batch dimension
+    image = Image.fromarray(image.astype('uint8'), 'RGB')
+    image = transform(image)
+    image = image.unsqueeze(0)  # Add batch dimension
+
+    # Run inference
     with torch.no_grad():
+        model.eval()
         outputs = model(image)
-        _, predicted = torch.max(outputs, 1)
-        return class_names[predicted.item()]
+        _, preds = torch.max(outputs, 1)
+        predicted_class = class_names[preds.item()]
+
+    return predicted_class
 
-# Create Gradio interface
+# Define the Gradio interface
 iface = gr.Interface(
-    fn=predict,
-    inputs=gr.inputs.Image(type="pil", label="Upload an Image"),
-    outputs=gr.outputs.Label(num_top_classes=4, label="Prediction"),
-    live=True
+    fn=predict, 
+    inputs=gr.inputs.Image(type="numpy", label="Upload Image"),
+    outputs=gr.outputs.Label(num_top_classes=1, label="Predicted Class"),
+    title="Monkeypox Classifier",
+    description="Upload an image of skin lesions to classify the disease."
 )
 
-# Launch the interface
+# Launch the app
 iface.launch()