Update app.py

app.py

@@ -1,28 +1,39 @@
+import gradio as gr
+import tensorflow as tf
 import requests
 import tempfile
-import tensorflow as tf
 
-# URL to the model file in your Hugging Face repository
-url = 'https://huggingface.co/datasets/yolac/BacterialMorphologyClassification'
-# Download the model file
-response = requests.get(url, stream=True)
+# Define the URL for your model and dataset
+model_url = 'https://huggingface.co/yolac/BacterialMorphologyClassification/resolve/main/bacterial_morphology_classification_model.h5'
 
-# Check if the download was successful
+# Download the model from the URL and load it into a temporary file
+response = requests.get(model_url, stream=True)
 if response.status_code == 200:
-    # Create a temporary file to save the model
-    with tempfile.NamedTemporaryFile(suffix=".h5", delete=False) as temp_file:
-        for chunk in response.iter_content(chunk_size=8192):
-            if chunk:
-                temp_file.write(chunk)
-        temp_file_path = temp_file.name  # Get the path to the temporary file
-
-    # Load the pre-trained model
-    try:
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.h5') as temp_file:
+        temp_file.write(response.content)
+        temp_file_path = temp_file.name
         model = tf.keras.models.load_model(temp_file_path)
-        print("Model loaded successfully.")
-    except OSError as e:
-        print(f"Error loading the model: {e}")
 else:
-    print("Failed to download the model. Status code:", response.status_code)
+    raise Exception(f"Failed to download the model. Status code: {response.status_code}")
+
+# Define the function for prediction
+def predict_image(img):
+    img = img.resize((224, 224))
+    img_array = tf.keras.preprocessing.image.img_to_array(img) / 255.0
+    img_array = tf.expand_dims(img_array, axis=0)
+    prediction = model.predict(img_array)
+    classes = ['cocci', 'bacilli', 'spirilla']
+    predicted_class = classes[prediction.argmax()]
+    return predicted_class
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict_image,
+    inputs=gr.inputs.Image(shape=(224, 224)),
+    outputs=gr.outputs.Label(num_top_classes=3),
+    title="Bacterial Morphology Classification",
+    description="Upload an image of bacterial morphology to classify it as cocci, bacilli, or spirilla."
+)
 
-# Now you can use the `model` object for predictions or further processing.
+# Launch the Gradio app
+iface.launch()