Update app.py

app.py

@@ -1,56 +1,100 @@
-import gradio as gr
 import torch
+import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
+import requests
+import gradio as gr
+import os
+
+# Define the model architecture
+class BacterialMorphologyClassifier(nn.Module):
+    def __init__(self):
+        super(BacterialMorphologyClassifier, self).__init__()
+        self.feature_extractor = nn.Sequential(
+            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+            nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2, stride=2),
+        )
+        self.fc = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(64 * 56 * 56, 128),
+            nn.ReLU(),
+            nn.Dropout(0.5),
+            nn.Linear(128, 3),
+            nn.Softmax(dim=1),
+        )
+    
+    def forward(self, x):
+        x = self.feature_extractor(x)
+        x = self.fc(x)
+        return x
 
 # Load the model
-def load_model():
-    model_path = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth"
-    model = torch.load(model_path, map_location=torch.device('cpu'))
-    model.eval()  # Set model to evaluation mode
-    return model
+MODEL_PATH = "model.pth"
+model = BacterialMorphologyClassifier()
 
-model = load_model()
+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')))
+    model.eval()
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Error loading the model: {e}")
 
-# Define class labels
-class_labels = {0: "Cocci", 1: "Bacilli", 2: "Spirilla"}
+# Define image preprocessing to match training preprocessing
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),  # Resize to match model input size
+    transforms.ToTensor(),  # Convert to a tensor
+    transforms.Normalize(mean=[0, 0, 0], std=[1/255, 1/255, 1/255]),  # Scale pixel values to [0, 1]
+])
 
 # Prediction function
-def predict_image(image):
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Lambda(lambda x: x / 255.0)  # Rescale pixel values to [0, 1]
-    ])
-    image_tensor = transform(image).unsqueeze(0)  # Add batch dimension
-
-    with torch.no_grad():
-        outputs = model(image_tensor)
-        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
-        predicted_class = probabilities.argmax().item()
-    return class_labels[predicted_class], {class_labels[i]: float(prob) for i, prob in enumerate(probabilities)}
-
-# Example images (these should be in the same directory as your app.py)
-example_images = ["https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20290.jpg", "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20565.jpg", "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%208.jpg"]
-
-# Create a Gradio interface
-def create_gradio_interface():
-    iface = gr.Interface(
-        fn=predict_image,
-        inputs=gr.inputs.Image(type="pil", label="Upload an image"),
-        outputs=[
-            gr.outputs.Label(num_top_classes=3, label="Predicted Class"),
-            gr.outputs.JSON(label="Class Probabilities")
-        ],
-        examples=example_images,
-        title="Bacterial Morphology Classification",
-        description="This app classifies bacterial morphology into **Cocci**, **Bacilli**, or **Spirilla** using a fine-tuned PyTorch model.",
-    )
-    return iface
+def predict(image):
+    try:
+        # Convert the image to a tensor
+        image_tensor = transform(image).unsqueeze(0)
+        
+        # Perform prediction
+        with torch.no_grad():  # Ensure no gradients are calculated
+            output = model(image_tensor)
+        
+        # Class mapping
+        class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
+        
+        # Return the predicted class and confidence
+        predicted_class = class_labels[output.argmax().item()]
+        confidence = output.max().item()  # Softmax value as confidence
+        return f"Predicted Class: {predicted_class}\nConfidence: {confidence:.2f}"
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Define example images
+examples = [
+    ["https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20290.jpg"],
+    ["https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20565.jpg"],
+    ["https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%208.jpg"],
+]
+
+# Set up Gradio interface
+interface = 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.",
+    examples=examples,
+)
 
 # Launch the app
 if __name__ == "__main__":
-    app = create_gradio_interface()
-    app.launch()
-
-    
+    interface.launch()