Update app.py

app.py

@@ -1,63 +1,67 @@
 import gradio as gr
 import torch
-import torch.nn as nn
-from torchvision import transforms
+import torchvision.transforms as transforms
 from PIL import Image
+from model import load_model
 
-# Define your model class
-class CustomModel(nn.Module):
-    def __init__(self):
-        super(CustomModel, self).__init__()
-        # Define your layers here (example)
-        self.fc = nn.Linear(512, 2)
-
-    def forward(self, x):
-        x = self.fc(x)
-        return x
-
-# Function to load your model
-def load_model(model_path):
-    model = CustomModel()
-    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
-    model.eval()
-    return model
-
-# Function to preprocess the image
-def preprocess_image(image):
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ])
+# 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(),
+])
+
+# Define the class names
+class_names = ['Normal', 'Monkeypox', 'Chickenpox', 'Measles']
+
+def predict_image(image):
+    # Preprocess the image
     image = Image.fromarray(image)
     image = transform(image).unsqueeze(0)
-    return image
-
-# Function to predict using the model
-def predict(image):
-    image = preprocess_image(image)
+    
+    # Perform inference
+    model.eval()
     with torch.no_grad():
-        output = model(image)
-    prediction = torch.argmax(output, dim=1).item()
-    class_names = ['Monkeypox', 'Healthy']
-    return class_names[prediction]
+        outputs = model(image)
+        _, predicted = torch.max(outputs, 1)
+        predicted_class = class_names[predicted.item()]
+    
+    return predicted_class
 
-# Load the model
-model = load_model('best_model_augmented.pth')
+# HTML for custom styling
+title = (
+    "<h1 style='text-align: center; font-family: Arial, sans-serif;'>Pox Classifier</h1>"
+    "<p style='text-align: center; font-family: Arial, sans-serif;'>"
+    "<a href='https://x.com/ronith_sharmila' target='_blank'>@your_twitter_handle</a>"
+    "</p>"
+)
 
-# Create the Gradio interface
-title = "<h1 style='text-align: center; margin-bottom: 10px;'>Pox Classifier</h1>"
-class_names_text = "<h3 style='text-align: center; margin-top: -10px;'>[Monkeypox, Healthy]</h3>"
-twitter_link = "<p style='text-align: center;'><a href='https://x.com/ronith_sharmila' target='_blank'>@your_twitter_handle</a></p>"
-support_link = "<p style='text-align: center;'><a href='https://paypal.me/ronithsharmila?country.x=US&locale.x=en_US' target='_blank'>Support the work <strong>here</strong></a></p>"
-disclaimer = "<p style='text-align: center; color: gray;'>Disclaimer: This model is for educational purposes only and should not be used for medical diagnosis.</p>"
+description = (
+    "<p style='text-align: center; font-family: Arial, sans-serif;'>"
+    "Upload an image of a skin lesion to classify it into one of the following categories:</p>"
+    "<p style='text-align: center; font-family: Arial, sans-serif; font-weight: bold;'>"
+    "Normal, Monkeypox, Chickenpox, Measles</p>"
+    "<p style='text-align: center; font-family: Arial, sans-serif;'>"
+    "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."
+    "</p>"
+    "<p style='text-align: center; font-family: Arial, sans-serif;'>"
+    "<a href='https://paypal.me/ronithsharmila?country.x=US&locale.x=en_US' target='_blank'>Support my work here</a>"
+    "</p>"
+)
 
+# Gradio app interface
 app = gr.Interface(
-    fn=predict,
+    fn=predict_image,
     inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=gr.Label(label="Prediction"),
+    outputs=gr.Textbox(label="Prediction"),
     title=title,
-    description=class_names_text + twitter_link + support_link + disclaimer
+    description=description,
+    theme="default"  # Use the default theme or specify a different theme if available
 )
 
+# Launch the app
 app.launch()