Update app.py

app.py

@@ -1,66 +1,42 @@
 import gradio as gr
-import tensorflow as tf
 from tensorflow.keras.models import load_model
-from tensorflow.keras.preprocessing.image import img_to_array
+from huggingface_hub import hf_hub_download
 import numpy as np
 from PIL import Image
-import os
-from huggingface_hub import hf_hub_download
 
-# Cache model locally
-MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.keras"
-
-# Download the model if not present
-if not os.path.exists(MODEL_PATH):
-    print("Downloading model...")
-    hf_hub_download(
-        repo_id="yolac/BacterialMorphologyClassification",
-        filename="model.keras",
-        local_dir="./",
-    )
+# Define constants
+MODEL_REPO = "yolac/BacterialMorphologyClassification"
+MODEL_FILENAME = "model.keras"
+MODEL_PATH = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILENAME)
 
 # Load the model
 print("Loading model...")
 model = load_model(MODEL_PATH)
 
-# Define class labels
-class_labels = ["Cocci", "Bacilli", "Spirilla"]
-
 # Preprocessing function
 def preprocess_image(image):
-    image = image.resize((224, 224))  # Resize to model input size
-    image = img_to_array(image)
-    image = np.expand_dims(image, axis=0)
-    image = image / 255.0  # Normalize to [0, 1]
-    return image
+    image = image.resize((224, 224))  # Adjust size as per your model input
+    image_array = np.array(image) / 255.0  # Normalize to [0, 1]
+    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
+    return image_array
 
 # Prediction function
-def classify_bacteria(image):
-    processed_image = preprocess_image(image)
-    predictions = model.predict(processed_image)
-    predicted_class = np.argmax(predictions, axis=-1)[0]
-    confidence = predictions[0][predicted_class]
-    return f"{class_labels[predicted_class]} (Confidence: {confidence:.2f})"
-
-# Gradio interface
-title = "Bacterial Morphology Classifier"
-description = (
-    "Upload an image of bacteria, and the model will classify it into one of three types: "
-    "**Cocci**, **Bacilli**, or **Spirilla**."
-)
-
+def predict(image):
+    image_array = preprocess_image(image)
+    predictions = model.predict(image_array)
+    class_names = ["Cocci", "Bacilli", "Spirilla"]
+    predicted_class = class_names[np.argmax(predictions)]
+    return f"Predicted Class: {predicted_class}"
+
+# Gradio Interface
 interface = gr.Interface(
-    fn=classify_bacteria,
+    fn=predict,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title=title,
-    description=description,
-    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"]
-    ],
+    title="Bacterial Morphology Classification",
+    description="Upload an image of bacteria to classify as Cocci, Bacilli, or Spirilla."
 )
 
+# Launch the app
 if __name__ == "__main__":
     interface.launch()