Update app.py

app.py

@@ -1,58 +1,63 @@
 import gradio as gr
 import torch
-import torchvision.transforms as transforms
+import torch.nn as nn
+from torchvision import transforms
 from PIL import Image
-from model import load_model
 
-# Load model
-model_path = 'best_model_augmented.pth'
-model = load_model(model_path)
+# 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)
 
-# Define preprocessing transformations
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-])
+    def forward(self, x):
+        x = self.fc(x)
+        return x
 
-# Define the class names
-class_names = ['Normal', 'Monkeypox', 'Chickenpox', 'Measles']
+# 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
 
-def predict_image(image):
-    # Preprocess the image
+# 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]),
+    ])
     image = Image.fromarray(image)
     image = transform(image).unsqueeze(0)
-    
-    # Perform inference
-    model.eval()
+    return image
+
+# Function to predict using the model
+def predict(image):
+    image = preprocess_image(image)
     with torch.no_grad():
-        outputs = model(image)
-        _, predicted = torch.max(outputs, 1)
-        predicted_class = class_names[predicted.item()]
-    
-    return predicted_class
-
-# HTML for custom styling
-title = "<h1 style='text-align: center; font-family: Arial, sans-serif;'>PoxNet</h1>"
-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>"
-)
+        output = model(image)
+    prediction = torch.argmax(output, dim=1).item()
+    class_names = ['Monkeypox', 'Healthy']
+    return class_names[prediction]
+
+# Load the model
+model = load_model('best_model_augmented.pth')
+
+# 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>"
 
-# Gradio app interface
 app = gr.Interface(
-    fn=predict_image,
+    fn=predict,
     inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=gr.Textbox(label="Prediction"),
+    outputs=gr.Label(label="Prediction"),
     title=title,
-    description=description,
-    theme="default"  # Use the default theme or specify a different theme if available
+    description=class_names_text + twitter_link + support_link + disclaimer
 )
 
-# Launch the app
 app.launch()