Update app.py

app.py

@@ -1,44 +1,44 @@
 import gradio as gr
 import torch
-from torchvision import transforms
+import torchvision.transforms as transforms
 from PIL import Image
-from model import CustomModel, load_model
+from model import load_model
 
-# Define the transformation
+# Load model
+model_path = 'best_model_augmented.pth'
+model = load_model(model_path)
+
+# Define preprocessing transformations
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
 ])
 
-# Load the model
-model = load_model('best_model_augmented.pth')
-
 # Define the class names
 class_names = ['Normal', 'Monkeypox', 'Chickenpox', 'Measles']
 
-def predict(image):
+def predict_image(image):
     # Preprocess the image
-    image = Image.fromarray(image.astype('uint8'), 'RGB')
-    image = transform(image)
-    image = image.unsqueeze(0)  # Add batch dimension
-
-    # Run inference
+    image = Image.fromarray(image)
+    image = transform(image).unsqueeze(0)
+    
+    # Perform inference
+    model.eval()
     with torch.no_grad():
-        model.eval()
         outputs = model(image)
-        _, preds = torch.max(outputs, 1)
-        predicted_class = class_names[preds.item()]
-
+        _, predicted = torch.max(outputs, 1)
+        predicted_class = class_names[predicted.item()]
+    
     return predicted_class
 
-# Define the Gradio interface using the updated API
-iface = gr.Interface(
-    fn=predict, 
+# Gradio app interface
+app = gr.Interface(
+    fn=predict_image,
     inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=gr.Label(num_top_classes=1, label="Predicted Class"),
-    title="Monkeypox Classifier",
-    description="Upload an image of skin lesions to classify the disease."
+    outputs=gr.Textbox(label="Prediction"),
+    title="Pox Classifier (Normal, Monkeypox, Chickenpox, Measles)",
+    description="Upload an image of a skin lesion to classify it into one of the following categories: Normal, Monkeypox, Chickenpox, or Measles. Please note that this model is not a substitute for professional medical advice, diagnosis, or treatment. Always seek the advice of your physician or other qualified health provider with any questions you may have regarding a medical condition."
 )
 
 # Launch the app
-iface.launch()
+app.launch()