Update app.py

app.py

@@ -1,101 +1,86 @@
+import streamlit as st
 import torch
-import torch.nn as nn
 from torchvision import transforms
 from PIL import Image
-import requests
-import gradio as gr
-import os
+import json
 
-# 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 Model
+@st.cache_resource
+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
 
-# Load the model
-MODEL_PATH = "model.pth"
-model = BacterialMorphologyClassifier()
+# Prediction Function
+def predict_image(model, image):
+    # Transform the 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
 
-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}")
+    with torch.no_grad():
+        outputs = model(image_tensor)
+        probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
+        predicted_class = probabilities.argmax().item()
+    return predicted_class, probabilities.numpy()
 
-# 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]
-])
+# Class Labels
+def get_class_labels():
+    # Define your class labels here
+    return {0: "Cocci", 1: "Bacilli", 2: "Spirilla"}
 
-# Prediction function
-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)}"
+# Streamlit App
+st.set_page_config(page_title="Bacterial Morphology Classification", page_icon="🦠")
 
-# 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"],
+st.title("🦠 Bacterial Morphology Classification")
+st.markdown(
+    "This app classifies bacterial morphology into **Cocci**, **Bacilli**, or **Spirilla** using a fine-tuned PyTorch model."
+)
+
+# Example Images
+st.subheader("Example Images")
+example_images = [
+    "https://huggingface.co/datasets/yolac/BacterialMorphologyClassification/resolve/main/img%20290.jpg",  # Replace with actual paths to example images
+    "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,
+for img_path in example_images:
+    img = Image.open(img_path).convert("RGB")
+    st.image(img, caption=f"Example Image: {img_path}", use_column_width=True)
+
+# File Upload
+uploaded_file = st.file_uploader("Upload a bacterial image for classification:", type=["jpg", "jpeg", "png"])
+
+if uploaded_file:
+    # Display the uploaded image
+    image = Image.open(uploaded_file).convert("RGB")
+    st.image(image, caption="Uploaded Image", use_column_width=True)
+
+    # Load Model and Predict
+    with st.spinner("Classifying..."):
+        model = load_model()
+        class_labels = get_class_labels()
+        predicted_class, probabilities = predict_image(model, image)
+        predicted_label = class_labels[predicted_class]
+
+        # Display Results
+        st.success(f"Prediction: **{predicted_label}**")
+        st.write("Class Probabilities:")
+        st.json({class_labels[i]: f"{prob:.2%}" for i, prob in enumerate(probabilities)})
+
+# Sidebar Info
+st.sidebar.title("Classifies bacterial images into cocci, bacilli, or spirilla")
+st.sidebar.markdown(
+    """
+    - **Author**: Yola Charara
+    - **Dataset**: [Bacterial Morphology Classification](https://huggingface.co/datasets/yolac/BacterialMorphologyClassification)
+    - **Model**: [MobileNetV2-based Classifier](https://huggingface.co/yolac/BacterialMorphologyClassification)
+    """
 )
 
-# Launch the app
-if __name__ == "__main__":
-    interface.launch()