Update app.py

app.py

@@ -1,78 +1,92 @@
 import torch
 import torch.nn as nn
-import gradio as gr
 from torchvision import transforms
 from PIL import Image
 import requests
-from io import BytesIO
+import gradio as gr
+import os
 
-# Define the PyTorch model architecture
-class MyPyTorchModel(nn.Module):
+# Define the model architecture
+class BacterialMorphologyClassifier(nn.Module):
     def __init__(self):
-        super(MyPyTorchModel, self).__init__()
-        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
-        self.relu = nn.ReLU()
-        self.flatten = nn.Flatten()
-        self.fc1 = nn.Linear(32 * 224 * 224, 3)  # Adjust output size for 3 classes
-
+        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.conv1(x)
-        x = self.relu(x)
-        x = self.flatten(x)
-        x = self.fc1(x)
+        x = self.feature_extractor(x)
+        x = self.fc(x)
         return x
 
 # Load the model
-model_path = 'https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.pth'
-model = MyPyTorchModel()
-model.load_state_dict(torch.load(model_path))
-model.eval()
+MODEL_PATH = "model.pth"
+model = BacterialMorphologyClassifier()
 
-# Define image transformations
+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()
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Error loading the model: {e}")
+
+# Define image preprocessing
 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]),
 ])
 
-# Define a function to predict the class of the image
+# Prediction function
 def predict(image):
-    image = Image.fromarray(image).convert('RGB')
-    image = transform(image).unsqueeze(0)
-
-    with torch.no_grad():
-        outputs = model(image)
-        _, predicted = torch.max(outputs, 1)
-
-    class_labels = ['cocci', 'bacilli', 'spirilla']
-    predicted_label = class_labels[predicted.item()]
-
-    return predicted_label
-
-# URLs for 3 example images
-example_image_urls = [
-    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20290.jpg",  # Replace with the actual URL
-    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20565.jpg",  # Replace with the actual URL
-    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%208.jpg",  # Replace with the actual URL
-]
-
-# Function to download and open example images
-def load_example_image(url):
-    response = requests.get(url)
-    image = Image.open(BytesIO(response.content))
-    return image
-
-# Load the example images
-example_images = [load_example_image(url) for url in example_image_urls]
+    try:
+        # Convert the image to a tensor
+        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}"
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-# Create a Gradio interface
-iface = gr.Interface(
+# Set up Gradio interface
+interface = gr.Interface(
     fn=predict,
-    inputs=gr.inputs.Image(shape=(224, 224), label="Upload an image or use examples"),
-    outputs="text",
-    title="Bacterial Morphology Classifier",
-    description="Classify images of bacteria into cocci, bacilli, or spirilla.",
-    examples=example_images  # Add example images to the Gradio app
+    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.",
 )
 
 # Launch the app
-iface.launch()
+if __name__ == "__main__":
+    interface.launch()