Update app.py

app.py

@@ -1,62 +1,72 @@
 import gradio as gr
 import tensorflow as tf
-from PIL import Image
+from tensorflow.keras.models import load_model
+from tensorflow.keras.preprocessing.image import img_to_array
 import numpy as np
-import logging
-
-# Set up logging for debugging
-logging.basicConfig(level=logging.DEBUG)
+from PIL import Image
+import os
 
-# Load the .keras model
+# Cache model during setup
 MODEL_PATH = "https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/model.keras"
-logging.debug("Starting model loading...")
-try:
-    model = tf.keras.models.load_model(MODEL_PATH)
-    logging.debug("Model loaded successfully.")
-except Exception as e:
-    logging.error(f"Error loading the model: {str(e)}")
-
-# Define image preprocessing transformations
+
+# Check if the model exists locally, otherwise download
+if not os.path.exists(MODEL_PATH):
+    from huggingface_hub import hf_hub_download
+    print("Downloading model...")
+    hf_hub_download(
+        repo_id="yolac/BacterialMorphologyClassification",
+        filename="model.keras",
+        local_dir="./",
+        local_dir_use_symlinks=False
+    )
+
+# 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):
-    logging.debug("Preprocessing image...")
-    image = image.resize((224, 224))  # Resize to match model input size
-    image_array = np.array(image) / 255.0  # Normalize pixel values
-    if len(image_array.shape) == 2:  # If grayscale, convert to RGB
-        image_array = np.stack([image_array] * 3, axis=-1)
-    image_array = np.expand_dims(image_array, axis=0)  # Add batch dimension
-    return image_array
-
-# Define the prediction function
-def predict(image):
-    try:
-        # Preprocess the image
-        image_array = preprocess_image(image)
-        logging.debug("Image preprocessing completed.")
-        
-        # Make prediction
-        predictions = model.predict(image_array)
-        prediction = np.argmax(predictions, axis=1)[0]
-        
-        # Class mapping
-        class_labels = {0: 'cocci', 1: 'bacilli', 2: 'spirilla'}
-        
-        # Log prediction details
-        logging.debug(f"Predicted class: {class_labels[prediction]}, Confidence: {predictions[0][prediction]}")
-        
-        # Return prediction result
-        return class_labels[prediction], float(predictions[0][prediction])
-    except Exception as e:
-        logging.error(f"Error during prediction: {str(e)}")
-        return "Error", 0.0
-
-# Create a Gradio interface
-gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil", label="Upload an image"),
-    outputs=["text", "number"],
+    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
+
+# Prediction function
+def classify_bacteria(image):
+    # Preprocess the input image
+    processed_image = preprocess_image(image)
+    # Get model predictions
+    predictions = model.predict(processed_image)
+    predicted_class = np.argmax(predictions, axis=-1)[0]
+    confidence = predictions[0][predicted_class]
+    # Return the class and confidence
+    return f"{class_labels[predicted_class]} (Confidence: {confidence:.2f})"
+
+# Define 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**."
+)
+
+# Gradio UI
+interface = gr.Interface(
+    fn=classify_bacteria,
+    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"]
-    ]
-).launch(server_name="0.0.0.0", server_port=7862, debug=True)
+    ],
+)
+
+# Launch the app
+if __name__ == "__main__":
+    interface.launch()