Update app.py

app.py

@@ -1,10 +1,11 @@
+import gradio as gr
 import torch
 import torch.nn as nn
-import gradio as gr
 from torchvision import transforms
 from PIL import Image
+import io
 
-# Define the model architecture
+# Define the model architecture that matches the saved .pth file
 class BacterialMorphologyClassifier(nn.Module):
     def __init__(self):
         super(BacterialMorphologyClassifier, self).__init__()
@@ -22,6 +23,7 @@ class BacterialMorphologyClassifier(nn.Module):
             nn.ReLU(),
             nn.Dropout(0.5),
             nn.Linear(128, 3),
+            nn.Softmax(dim=1),
         )
     
     def forward(self, x):
@@ -29,17 +31,11 @@ class BacterialMorphologyClassifier(nn.Module):
         x = self.fc(x)
         return x
 
-# Load the model
-MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
+# Load the model and weights
 model = BacterialMorphologyClassifier()
-try:
-    # 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 model: {e}")
-    raise e
+MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
+state_dict = torch.hub.load_state_dict_from_url(MODEL_PATH, map_location=torch.device('cpu'))
+model.load_state_dict(state_dict, strict=False)
 model.eval()
 
 # Define image preprocessing transformations
@@ -49,41 +45,36 @@ transform = transforms.Compose([
     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
+# Define Gradio interface
 def predict(image):
     try:
         # Preprocess the image
         image_tensor = transform(image).unsqueeze(0)
-
-        # 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}
+        
+        # Make prediction
+        output = model(image_tensor)
+        prediction = output.argmax().item()
+        
+        # Class mapping
+        class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
+        
+        # Return prediction result
+        return class_labels[prediction], output.max().item()
     except Exception as e:
         return {'error': str(e)}
 
-# Example input images (provide paths or URLs)
-example_images = [
-    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer?row=0&image-viewer=52B421CB70A43313B278D5DD2C58CECE56343012",  # Replace with the actual paths to your example images
-    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer/default/train?p=2&row=201&image-viewer=558EA847F2267CECF4E2CFF6352F9D8888E9A72F",
-    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer/default/train?p=2&row=201&image-viewer=8FBAF2C52C256A392660811C5659788734821C3A"
-]
-
-# Set up Gradio interface with examples
+# Create 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=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'.",
-    examples=example_images  # Provide the example image paths
+    inputs=gr.Image(type="pil", label="Upload an image"),
+    outputs=[gr.Label(num_top_classes=3, label="Predicted Class"), gr.Number(label="Confidence")],
+    examples=[
+        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer/default/train?p=2&row=201&image-viewer=8FBAF2C52C256A392660811C5659788734821C3A",
+        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer/default/train?p=2&image-viewer=AEF1AA2978EEB77362DA9CCC8792473666F7CDC6",
+        "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/viewer/default/train?image-viewer=C98E6CFAB26ECC3808C63185F6CCE90DE4E7C442"
+    ]
 )
 
 # Launch the app
-iface.launch(server_name="0.0.0.0", server_port=5000, share=True)
+if __name__ == "__main__":
+    iface.launch()