Update app.py

app.py

@@ -1,10 +1,8 @@
 import torch
 import torch.nn as nn
+import gradio as gr
 from torchvision import transforms
 from PIL import Image
-import requests
-import gradio as gr
-import os
 
 # Define the model architecture
 class BacterialMorphologyClassifier(nn.Module):
@@ -24,7 +22,6 @@ class BacterialMorphologyClassifier(nn.Module):
             nn.ReLU(),
             nn.Dropout(0.5),
             nn.Linear(128, 3),
-            nn.Softmax(dim=1),
         )
     
     def forward(self, x):
@@ -33,60 +30,52 @@ class BacterialMorphologyClassifier(nn.Module):
         return x
 
 # Load the model
-MODEL_PATH = "model.pth"
+MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
 model = BacterialMorphologyClassifier()
-
 try:
-    # Download the model if it doesn't exist
-    if not os.path.exists(MODEL_PATH):
-        print("Downloading the model...")
-        url = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
-        response = requests.get(url)
-        with open(MODEL_PATH, "wb") as f:
-            f.write(response.content)
-    # Load the model weights
-    model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')), strict=False)
-    model.eval()
+    # Download and load model state_dict
+    state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
+    model.load_state_dict(state_dict, strict=False)
     print("Model loaded successfully.")
 except Exception as e:
-    print(f"Error loading the model: {e}")
+    print(f"Error loading model: {e}")
+    raise e
+model.eval()
 
-# Define image preprocessing
+# Define image preprocessing transformations
 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]),
 ])
 
+# Class labels
+class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
+
 # Prediction function
 def predict(image):
     try:
-        # Convert the image to a tensor
+        # Preprocess the image
         image_tensor = transform(image).unsqueeze(0)
-        
-        # Perform prediction
-        output = model(image_tensor)
-        prediction = output.argmax().item()
-        
-        # Class mapping
-        class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
-        
-        # Return the predicted class and confidence
-        predicted_class = class_labels[prediction]
-        confidence = output.max().item()
-        return f"Predicted Class: {predicted_class}\nConfidence: {confidence:.2f}"
+
+        # Perform inference
+        with torch.no_grad():
+            output = model(image_tensor)
+            prediction = output.argmax().item()
+            confidence = torch.nn.functional.softmax(output, dim=1).max().item()
+
+        return {class_labels[prediction]: confidence}
     except Exception as e:
-        return f"Error: {str(e)}"
+        return {'error': str(e)}
 
 # Set up Gradio interface
-interface = gr.Interface(
+iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Text(label="Prediction"),
-    title="Bacterial Morphology Classification",
-    description="Upload an image of bacteria to classify it as cocci, bacilli, or spirilla.",
+    inputs=gr.inputs.Image(type="pil", label="Upload an image"),
+    outputs=gr.outputs.Label(num_top_classes=3, label="Predicted class and confidence"),
+    title="Bacterial Morphology Classifier",
+    description="Upload an image of a bacterial sample to classify it as 'cocci', 'bacilli', or 'spirilla'."
 )
 
 # Launch the app
-if __name__ == "__main__":
-    interface.launch()
+iface.launch(server_name="0.0.0.0", server_port=5000, share=True)