Update app.py

app.py

@@ -1,20 +1,22 @@
-import streamlit as st
+import gradio as gr
 import torch
 from torchvision import transforms
 from PIL import Image
-import json
 
-# Load Model
-@st.cache_resource
+# 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
 
-# Prediction Function
-def predict_image(model, image):
-    # Transform the image
+model = load_model()
+
+# Define class labels
+class_labels = {0: "Cocci", 1: "Bacilli", 2: "Spirilla"}
+
+# Prediction function
+def predict_image(image):
     transform = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
@@ -26,64 +28,29 @@ def predict_image(model, image):
         outputs = model(image_tensor)
         probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
         predicted_class = probabilities.argmax().item()
-    return predicted_class, probabilities.numpy()
-
-# Class Labels
-def get_class_labels():
-    # Define your class labels here
-    return {0: "Cocci", 1: "Bacilli", 2: "Spirilla"}
-
-# Streamlit App
-st.set_page_config(page_title="Bacterial Morphology Classification", page_icon="🦠")
-
-st.title("🦠 Bacterial Morphology Classification")
-st.markdown(
-    "This app classifies bacterial morphology into **Cocci**, **Bacilli**, or **Spirilla** using a fine-tuned PyTorch model."
-)
-
-# Display example images from local files
-st.subheader("Example Images")
-example_image_paths = [
-    "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"
-]
-
-for img_path in example_image_paths:
-    try:
-        img = Image.open(img_path).convert("RGB")
-        st.image(img, caption=f"Example Image: {img_path}", use_column_width=True)
-    except FileNotFoundError:
-        st.error(f"Example image {img_path} not found. Please ensure the image exists in the app directory.")
-
-# 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)
-    """
-)
+    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
+
+# Launch the app
+if __name__ == "__main__":
+    app = create_gradio_interface()
+    app.launch()